From a5b465b5884bbeccd825c68368d1945c937c7a69 Mon Sep 17 00:00:00 2001 From: Kent Date: Thu, 25 Jun 2026 00:13:36 +0800 Subject: [PATCH 001/187] feat(bedrock): add helper to neutralize orphaned tool blocks --- .../bedrock/chat/converse_transformation.py | 57 +++++++ .../chat/test_converse_transformation.py | 154 ++++++++++++++++++ 2 files changed, 211 insertions(+) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index a700c07d87a..4929b4a686d 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -43,6 +43,7 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAnnotation, ChatCompletionAssistantMessage, + ChatCompletionAssistantToolCall, ChatCompletionRedactedThinkingBlock, ChatCompletionResponseMessage, ChatCompletionSystemMessage, @@ -191,6 +192,62 @@ class AmazonConverseConfig(BaseConfig): return messages_copy + @staticmethod + def _has_orphaned_tool_blocks(messages: List[AllMessageValues]) -> bool: + return any( + (m.get("role") == "assistant" and m.get("tool_calls")) + or m.get("role") in ("tool", "function") + for m in messages + ) + + @staticmethod + def _neutralize_orphaned_tool_blocks( + messages: List[AllMessageValues], optional_params: dict + ) -> List[AllMessageValues]: + if optional_params.get( + "tools" + ) or not AmazonConverseConfig._has_orphaned_tool_blocks(messages): + return messages + + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, + ) + + def _tool_call_text(tool_call: ChatCompletionAssistantToolCall) -> str: + function = tool_call.get("function") or {} + name = function.get("name") or "unknown_tool" + arguments = function.get("arguments") or "" + return f"[tool call: {name}({arguments})]" + + def _result_text(message: AllMessageValues) -> str: + rendered = convert_content_list_to_str(message).strip() + return rendered or "" + + def _rewrite(message: AllMessageValues) -> AllMessageValues: + role = message.get("role") + if role == "assistant" and message.get("tool_calls"): + base_text = convert_content_list_to_str(message) + call_texts = [_tool_call_text(call) for call in message["tool_calls"]] + text = "\n".join(filter(None, [base_text, *call_texts])) + return ChatCompletionAssistantMessage(role="assistant", content=text) + if role in ("tool", "function"): + tool_call_id = message.get("tool_call_id") + name = message.get("name") + label = f"tool result for {tool_call_id or name or 'unknown'}" + return ChatCompletionUserMessage( + role="user", + content=f"[{label}: {_result_text(message)}]", + ) + return message + + verbose_logger.warning( + "litellm.bedrock: request has tool blocks in message history but no " + "`tools=` param; neutralizing orphaned tool blocks to text so Bedrock " + "accepts the request without a toolConfig. Non-text tool-result " + "payloads are dropped. Pass `tools=` to preserve structured tool calling." + ) + return [_rewrite(message) for message in messages] + @classmethod def get_config(cls): return { diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index d940f9f47a6..383a8417aaf 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5428,3 +5428,157 @@ async def test_grounding_source_and_query_rendered_as_text(): user_content = result[0]["content"] assert {"text": "Tokyo is the capital of Japan."} in user_content assert {"text": "What is the capital of Japan?"} in user_content + + +def _orphaned_tool_history_messages(): + return [ + {"role": "user", "content": "What's the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_abc", + "content": "Sunny, 25C", + }, + {"role": "user", "content": "Summarize our conversation so far."}, + ] + + +def test_neutralize_orphaned_tool_blocks_rewrites_when_no_tools(): + """No tools= but history has tool blocks: assistant tool_calls and the tool + result must be rewritten to text, with the structured tool fields gone and + tool_call_id preserved, so Bedrock accepts the request without a toolConfig + (#24158, #27138).""" + messages = _orphaned_tool_history_messages() + + result = AmazonConverseConfig._neutralize_orphaned_tool_blocks( + messages, optional_params={} + ) + + serialized = json.dumps(result) + assert "tool_calls" not in serialized + assert not any(m.get("role") in ("tool", "function") for m in result) + assert "get_weather" in serialized + # The arguments string contains quotes; after json.dumps the literal + # '{"city": "Paris"}' is escaped, so assert on quote-free tokens that survive. + assert "city" in serialized and "Paris" in serialized + assert "Sunny, 25C" in serialized + assert "call_abc" in serialized # tool_call_id correlation preserved + + +@pytest.mark.parametrize("tools_value", [[], None]) +def test_neutralize_orphaned_tool_blocks_rewrites_when_tools_empty(tools_value): + """tools=[] and tools=None are 'no usable tools'; the gate must be on + truthiness, not key presence, or these slip through and still emit + structured tool blocks with no toolConfig.""" + messages = _orphaned_tool_history_messages() + + result = AmazonConverseConfig._neutralize_orphaned_tool_blocks( + messages, optional_params={"tools": tools_value} + ) + + serialized = json.dumps(result) + assert "tool_calls" not in serialized + assert "get_weather" in serialized + + +def test_neutralize_orphaned_tool_blocks_rewrites_tool_result_only_history(): + """A role:"tool"-only history (no assistant tool_calls) must also be + neutralized; has_tool_call_blocks misses this, but the factory still emits a + lone toolResult with no toolConfig.""" + messages = [ + {"role": "user", "content": "hi"}, + {"role": "tool", "tool_call_id": "call_xyz", "content": "lookup result"}, + ] + + result = AmazonConverseConfig._neutralize_orphaned_tool_blocks( + messages, optional_params={} + ) + + assert not any(m.get("role") in ("tool", "function") for m in result) + serialized = json.dumps(result) + assert "lookup result" in serialized + assert "call_xyz" in serialized + + +def test_neutralize_orphaned_tool_blocks_non_text_result_marked_not_empty(): + """Non-text tool-result payloads (image/file) collapse to an explicit + marker, never an empty string (Bedrock rejects empty text blocks) and never + a silent drop.""" + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "c1", "type": "function", + "function": {"name": "render", "arguments": "{}"}} + ], + }, + { + "role": "tool", + "tool_call_id": "c1", + "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}], + }, + ] + + result = AmazonConverseConfig._neutralize_orphaned_tool_blocks( + messages, optional_params={} + ) + + rewritten = next(m for m in result if m.get("role") == "user" and m is not messages[0]) + text = rewritten["content"] + assert text.strip() # never empty + assert "non-text tool result omitted" in text + + +def test_neutralize_orphaned_tool_blocks_noop_when_tools_present(): + """When a non-empty tools= is provided, tool blocks are legitimate and must + be left untouched (returns the same object, no rewriting).""" + messages = _orphaned_tool_history_messages() + + result = AmazonConverseConfig._neutralize_orphaned_tool_blocks( + messages, + optional_params={"tools": [{"type": "function", "function": {"name": "x"}}]}, + ) + + assert result is messages + + +def test_neutralize_orphaned_tool_blocks_noop_when_no_tool_history(): + """Plain conversation with no tool blocks is returned unchanged.""" + messages = [{"role": "user", "content": "hi"}] + + result = AmazonConverseConfig._neutralize_orphaned_tool_blocks( + messages, optional_params={} + ) + + assert result is messages + + +def test_neutralize_orphaned_tool_blocks_logs_warning(caplog): + """Neutralization must surface at WARNING level so a developer who forgot + tools= sees it instead of a silent degrade.""" + messages = _orphaned_tool_history_messages() + + with caplog.at_level("WARNING"): + AmazonConverseConfig._neutralize_orphaned_tool_blocks( + messages, optional_params={} + ) + + assert any( + "neutralizing orphaned tool blocks" in record.getMessage() + for record in caplog.records + ) From 600924af329bf67ff29d0d4c19b1b63eb6e4839a Mon Sep 17 00:00:00 2001 From: Kent Date: Thu, 25 Jun 2026 00:24:58 +0800 Subject: [PATCH 002/187] fix(bedrock): neutralize orphaned tool blocks instead of raising or injecting dummy tool (#24158, #27138) --- .../bedrock/chat/converse_transformation.py | 26 +-- tests/local_testing/test_function_calling.py | 10 +- .../chat/test_converse_transformation.py | 174 ++++++++++++++++++ 3 files changed, 183 insertions(+), 27 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 4929b4a686d..79d5ac51fa8 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -66,9 +66,7 @@ from litellm.types.utils import ( Usage, ) from litellm.utils import ( - add_dummy_tool, any_assistant_message_has_thinking_blocks, - has_tool_call_blocks, last_assistant_with_tool_calls_has_no_thinking_blocks, supports_reasoning, token_counter, @@ -1647,26 +1645,6 @@ class AmazonConverseConfig(BaseConfig): headers: Optional[dict] = None, drop_params: bool = False, ) -> CommonRequestObject: - ## VALIDATE REQUEST - """ - Bedrock doesn't support tool calling without `tools=` param specified. - """ - if ( - "tools" not in optional_params - and messages is not None - and has_tool_call_blocks(messages) - ): - if litellm.modify_params: - optional_params["tools"] = add_dummy_tool( - custom_llm_provider="bedrock_converse" - ) - else: - raise litellm.UnsupportedParamsError( - message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.", - model="", - llm_provider="bedrock", - ) - # Drop thinking param if thinking is enabled but thinking_blocks are missing # This prevents the error: "Expected thinking or redacted_thinking, but found tool_use" # @@ -1763,6 +1741,8 @@ class AmazonConverseConfig(BaseConfig): messages, model=model ) + messages = self._neutralize_orphaned_tool_blocks(messages, optional_params) + # Convert last user message to guarded_text if guardrailConfig is present messages = self._convert_consecutive_user_messages_to_guarded_text( messages, optional_params @@ -1822,6 +1802,8 @@ class AmazonConverseConfig(BaseConfig): messages, model=model ) + messages = self._neutralize_orphaned_tool_blocks(messages, optional_params) + # Convert last user message to guarded_text if guardrailConfig is present messages = self._convert_consecutive_user_messages_to_guarded_text( messages, optional_params diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index 3c7e004b62e..b51da95c4f0 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -301,11 +301,11 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [ @pytest.mark.parametrize( "model, messages, expect_unsupported_params_error", [ - # Bedrock Converse still requires modify_params to inject the dummy tool. + # Bedrock Converse neutralizes orphaned tool blocks to text; no error. ( "anthropic.claude-3-sonnet-20240229-v1:0", _PARALLEL_TOOL_HISTORY_MESSAGES, - True, + False, ), # Anthropic Messages API: dummy tool is injected without modify_params. ( @@ -341,12 +341,12 @@ def test_parallel_function_call_anthropic_error_msg( """ Tool history without an explicit ``tools`` param: - - Bedrock **Converse** still raises ``UnsupportedParamsError`` unless - ``litellm.modify_params`` is enabled (dummy tool is only added there). + - Bedrock **Converse** neutralizes the orphaned tool blocks into plain text + and sends no ``toolConfig`` (see #24158, #27138). It no longer raises. - **Anthropic** (and Bedrock Invoke via ``AnthropicConfig.transform_request``) always get a dummy tool so CLIs work with ``modify_params`` left off. - Reference Issue: https://github.com/BerriAI/litellm/issues/5747, https://github.com/BerriAI/litellm/issues/5388 + Reference Issue: https://github.com/BerriAI/litellm/issues/24158, https://github.com/BerriAI/litellm/issues/27138 """ # Ensure modify_params is False so Bedrock Converse path still raises. # (other tests in this file set it to True and don't reset it) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 383a8417aaf..8750d587dda 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5582,3 +5582,177 @@ def test_neutralize_orphaned_tool_blocks_logs_warning(caplog): "neutralizing orphaned tool blocks" in record.getMessage() for record in caplog.records ) + + +def _assert_no_structured_tool_blocks(result): + """A valid Bedrock body for a neutralized request has no tool config AND no + structured tool blocks in messages. Checking only toolConfig is insufficient: + deleting the raise without rewriting still leaves toolUse/toolResult, the + exact shape Bedrock rejects.""" + assert "toolConfig" not in result + serialized = json.dumps(result) + assert "toolUse" not in serialized + assert "toolResult" not in serialized + + +def test_transform_request_no_tools_with_tool_history_succeeds_24158(monkeypatch): + """#24158: a compaction-style call (tool blocks in history, no tools=) must + not raise and must send no toolConfig or structured tool blocks, on + default settings.""" + monkeypatch.setattr(litellm, "modify_params", False) + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-opus-4-5-20251101-v1:0", + messages=_orphaned_tool_history_messages(), + optional_params={}, + litellm_params={}, + headers={}, + ) + + _assert_no_structured_tool_blocks(result) + serialized = json.dumps(result) + assert "get_weather" in serialized + assert "Sunny, 25C" in serialized + + +def test_transform_request_tool_unsupported_model_no_toolconfig_27138(monkeypatch): + """#27138: a tool-incapable model with tool blocks in history and no tools= + must not get a toolConfig/toolUse/toolResult injected (which Bedrock would + 400 on), even with modify_params on.""" + monkeypatch.setattr(litellm, "modify_params", True) + config = AmazonConverseConfig() + + result = config.transform_request( + model="meta.llama3-2-3b-instruct-v1:0", + messages=_orphaned_tool_history_messages(), + optional_params={}, + litellm_params={}, + headers={}, + ) + + _assert_no_structured_tool_blocks(result) + + +@pytest.mark.parametrize("tools_value", [[], None]) +def test_transform_request_empty_tools_with_tool_history(monkeypatch, tools_value): + """tools=[] / tools=None must be neutralized like no tools at all; a + key-presence gate would skip them and emit toolUse/toolResult with no + toolConfig.""" + monkeypatch.setattr(litellm, "modify_params", False) + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-opus-4-5-20251101-v1:0", + messages=_orphaned_tool_history_messages(), + optional_params={"tools": tools_value}, + litellm_params={}, + headers={}, + ) + + _assert_no_structured_tool_blocks(result) + + +def test_transform_request_tool_result_only_history(monkeypatch): + """A role:"tool"-only history (no assistant tool_calls) currently emits a + lone toolResult with no toolConfig; it must be neutralized.""" + monkeypatch.setattr(litellm, "modify_params", False) + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-opus-4-5-20251101-v1:0", + messages=[ + {"role": "user", "content": "hi"}, + {"role": "tool", "tool_call_id": "call_xyz", "content": "lookup result"}, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + _assert_no_structured_tool_blocks(result) + assert "lookup result" in json.dumps(result) + + +def test_transform_request_neutralized_tool_output_is_guarded(monkeypatch): + """With guardrailConfig present, a neutralized tool result that becomes the + trailing user turn must be emitted as guardContent, not plain text, so + untrusted tool output does not bypass the guardrail (neutralize must run + before guarded-text conversion).""" + monkeypatch.setattr(litellm, "modify_params", False) + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-opus-4-5-20251101-v1:0", + messages=[ + {"role": "user", "content": "look it up"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "secret tool output"}, + ], + optional_params={ + "guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "1"} + }, + litellm_params={}, + headers={}, + ) + + _assert_no_structured_tool_blocks(result) + serialized = json.dumps(result) + assert "guardContent" in serialized + assert "secret tool output" in serialized + + +@pytest.mark.asyncio +async def test_async_transform_request_no_tools_with_tool_history(monkeypatch): + """Async is a separate request assembler; it must neutralize identically.""" + monkeypatch.setattr(litellm, "modify_params", False) + config = AmazonConverseConfig() + + result = await config._async_transform_request( + model="us.anthropic.claude-opus-4-5-20251101-v1:0", + messages=_orphaned_tool_history_messages(), + optional_params={}, + litellm_params={}, + headers={}, + ) + + _assert_no_structured_tool_blocks(result) + assert "get_weather" in json.dumps(result) + + +def test_transform_request_with_tools_still_builds_toolconfig(monkeypatch): + """Guard: when a non-empty tools= IS provided, tool blocks are legitimate and + a toolConfig must still be produced (neutralization must not regress this).""" + monkeypatch.setattr(litellm, "modify_params", False) + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-opus-4-5-20251101-v1:0", + messages=_orphaned_tool_history_messages(), + optional_params={ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + }, + litellm_params={}, + headers={}, + ) + + assert "toolConfig" in result From edb7ec230ee6c42a5ccbbdbef53abf96e8d69d3f Mon Sep 17 00:00:00 2001 From: Kent Date: Thu, 25 Jun 2026 00:40:34 +0800 Subject: [PATCH 003/187] chore(bedrock): fix stale test comment; modernize types and narrow tool_calls for lint Fix the now-inaccurate comment in test_parallel_function_call_anthropic_error_msg (Bedrock Converse no longer raises; modify_params is forced off as a clean baseline that exercises the Anthropic dummy-tool path). Switch the new helper's annotations to builtin list[...] to satisfy the ruff-strict UP006 ceiling, and bind message.get("tool_calls") to a local so basedpyright narrows the union instead of tripping reportGeneralTypeIssues. Both keep the lint budgets ratcheting down rather than bumping ceilings. --- litellm/llms/bedrock/chat/converse_transformation.py | 11 ++++++----- tests/local_testing/test_function_calling.py | 3 ++- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 79d5ac51fa8..5781b97ce00 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -191,7 +191,7 @@ class AmazonConverseConfig(BaseConfig): return messages_copy @staticmethod - def _has_orphaned_tool_blocks(messages: List[AllMessageValues]) -> bool: + def _has_orphaned_tool_blocks(messages: list[AllMessageValues]) -> bool: return any( (m.get("role") == "assistant" and m.get("tool_calls")) or m.get("role") in ("tool", "function") @@ -200,8 +200,8 @@ class AmazonConverseConfig(BaseConfig): @staticmethod def _neutralize_orphaned_tool_blocks( - messages: List[AllMessageValues], optional_params: dict - ) -> List[AllMessageValues]: + messages: list[AllMessageValues], optional_params: dict + ) -> list[AllMessageValues]: if optional_params.get( "tools" ) or not AmazonConverseConfig._has_orphaned_tool_blocks(messages): @@ -223,9 +223,10 @@ class AmazonConverseConfig(BaseConfig): def _rewrite(message: AllMessageValues) -> AllMessageValues: role = message.get("role") - if role == "assistant" and message.get("tool_calls"): + tool_calls = message.get("tool_calls") + if role == "assistant" and tool_calls: base_text = convert_content_list_to_str(message) - call_texts = [_tool_call_text(call) for call in message["tool_calls"]] + call_texts = [_tool_call_text(call) for call in tool_calls] text = "\n".join(filter(None, [base_text, *call_texts])) return ChatCompletionAssistantMessage(role="assistant", content=text) if role in ("tool", "function"): diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index b51da95c4f0..e0984cd4311 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -348,7 +348,8 @@ def test_parallel_function_call_anthropic_error_msg( Reference Issue: https://github.com/BerriAI/litellm/issues/24158, https://github.com/BerriAI/litellm/issues/27138 """ - # Ensure modify_params is False so Bedrock Converse path still raises. + # Force modify_params off as a clean baseline: it exercises the Anthropic + # dummy-tool path, which injects regardless of modify_params # (other tests in this file set it to True and don't reset it) original_modify_params = litellm.modify_params litellm.modify_params = False From 989f5e06f5858aaec59f039e4f206b46a5d3ca5a Mon Sep 17 00:00:00 2001 From: Kent Date: Fri, 26 Jun 2026 12:34:06 +0800 Subject: [PATCH 004/187] feat(bedrock): add opt-out flag for orphaned-tool-block neutralization Address Greptile P1 review feedback on PR #31400. P1 (backward compat): gate neutralization behind a new litellm.bedrock_neutralize_orphaned_tool_blocks flag (default True, so both bugs stay fixed). Setting it False restores the legacy contract: raise UnsupportedParamsError, or inject a dummy tool under modify_params. A new _handle_orphaned_tool_blocks dispatcher selects the path; the pure _neutralize_orphaned_tool_blocks helper is unchanged. Adds tests for flag-off-raises, flag-off-with-modify_params-injects, and default-on. P1 (CI credentials): the live-network test_parallel_function_call_anthropic_error_msg no longer needs Bedrock creds for the Converse case. That case is dropped (its no-raise behavior is already covered offline in test_converse_transformation.py); the now-unused expect_unsupported_params_error param and dead pytest.raises branch are removed. --- litellm/__init__.py | 4 ++ .../bedrock/chat/converse_transformation.py | 30 +++++++- tests/local_testing/test_function_calling.py | 65 +++++------------ .../chat/test_converse_transformation.py | 71 +++++++++++++++++-- 4 files changed, 118 insertions(+), 52 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 5ae5942f32f..4b10130326e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -242,6 +242,10 @@ telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False)) modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) +bedrock_neutralize_orphaned_tool_blocks = ( + os.getenv("LITELLM_BEDROCK_NEUTRALIZE_ORPHANED_TOOL_BLOCKS", "true").lower() + == "true" +) use_chat_completions_url_for_anthropic_messages: bool = bool( os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False) ) # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 5781b97ce00..4bdba492ee5 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -66,7 +66,9 @@ from litellm.types.utils import ( Usage, ) from litellm.utils import ( + add_dummy_tool, any_assistant_message_has_thinking_blocks, + has_tool_call_blocks, last_assistant_with_tool_calls_has_no_thinking_blocks, supports_reasoning, token_counter, @@ -247,6 +249,30 @@ class AmazonConverseConfig(BaseConfig): ) return [_rewrite(message) for message in messages] + @staticmethod + def _handle_orphaned_tool_blocks( + messages: list[AllMessageValues], optional_params: dict + ) -> list[AllMessageValues]: + if litellm.bedrock_neutralize_orphaned_tool_blocks: + return AmazonConverseConfig._neutralize_orphaned_tool_blocks( + messages, optional_params + ) + + if "tools" in optional_params or not has_tool_call_blocks(messages): + return messages + + if litellm.modify_params: + optional_params["tools"] = add_dummy_tool( + custom_llm_provider="bedrock_converse" + ) + return messages + + raise litellm.utils.UnsupportedParamsError( + message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.", + model="", + llm_provider="bedrock", + ) + @classmethod def get_config(cls): return { @@ -1742,7 +1768,7 @@ class AmazonConverseConfig(BaseConfig): messages, model=model ) - messages = self._neutralize_orphaned_tool_blocks(messages, optional_params) + messages = self._handle_orphaned_tool_blocks(messages, optional_params) # Convert last user message to guarded_text if guardrailConfig is present messages = self._convert_consecutive_user_messages_to_guarded_text( @@ -1803,7 +1829,7 @@ class AmazonConverseConfig(BaseConfig): messages, model=model ) - messages = self._neutralize_orphaned_tool_blocks(messages, optional_params) + messages = self._handle_orphaned_tool_blocks(messages, optional_params) # Convert last user message to guarded_text if guardrailConfig is present messages = self._convert_consecutive_user_messages_to_guarded_text( diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index e0984cd4311..3d32e2e24f9 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -267,7 +267,6 @@ def test_aaparallel_function_call_with_anthropic_thinking(model): from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message - _PARALLEL_TOOL_HISTORY_MESSAGES = [ { "role": "user", @@ -299,20 +298,11 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [ @pytest.mark.parametrize( - "model, messages, expect_unsupported_params_error", + "model, messages", [ - # Bedrock Converse neutralizes orphaned tool blocks to text; no error. - ( - "anthropic.claude-3-sonnet-20240229-v1:0", - _PARALLEL_TOOL_HISTORY_MESSAGES, - False, - ), - # Anthropic Messages API: dummy tool is injected without modify_params. - ( - "claude-haiku-4-5-20251001", - _PARALLEL_TOOL_HISTORY_MESSAGES, - False, - ), + # Anthropic Messages API: a dummy tool is injected without modify_params, + # so tool history with no tools= completes instead of raising. + ("claude-haiku-4-5-20251001", _PARALLEL_TOOL_HISTORY_MESSAGES), ( "anthropic.claude-3-sonnet-20240229-v1:0", [ @@ -321,7 +311,6 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [ "content": "What's the weather like in San Francisco, Tokyo, and Paris? - give me 3 responses", } ], - False, ), ( "claude-haiku-4-5-20251001", @@ -331,22 +320,18 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [ "content": "What's the weather like in San Francisco, Tokyo, and Paris? - give me 3 responses", } ], - False, ), ], ) -def test_parallel_function_call_anthropic_error_msg( - model, messages, expect_unsupported_params_error -): +def test_parallel_function_call_anthropic_error_msg(model, messages): """ - Tool history without an explicit ``tools`` param: + Tool history without an explicit ``tools`` param must complete, not raise. - - Bedrock **Converse** neutralizes the orphaned tool blocks into plain text - and sends no ``toolConfig`` (see #24158, #27138). It no longer raises. - - **Anthropic** (and Bedrock Invoke via ``AnthropicConfig.transform_request``) - always get a dummy tool so CLIs work with ``modify_params`` left off. - - Reference Issue: https://github.com/BerriAI/litellm/issues/24158, https://github.com/BerriAI/litellm/issues/27138 + Anthropic (and Bedrock Invoke via ``AnthropicConfig.transform_request``) + inject a dummy tool so CLIs work with ``modify_params`` left off. Bedrock + Converse's no-raise behavior is covered offline in + ``tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py`` + (see #24158, #27138), which needs no live credentials. """ # Force modify_params off as a clean baseline: it exercises the Anthropic # dummy-tool path, which injects regardless of modify_params @@ -355,26 +340,14 @@ def test_parallel_function_call_anthropic_error_msg( litellm.modify_params = False try: litellm.set_verbose = True - - if expect_unsupported_params_error: - with pytest.raises(litellm.UnsupportedParamsError) as e: - second_response = litellm.completion( - model=model, - messages=messages, - temperature=0.2, - seed=22, - drop_params=True, - ) # get a new response from the model where it can see the function response - print("second response\n", second_response) - else: - second_response = litellm.completion( - model=model, - messages=messages, - temperature=0.2, - seed=22, - drop_params=True, - ) # get a new response from the model where it can see the function response - print("second response\n", second_response) + second_response = litellm.completion( + model=model, + messages=messages, + temperature=0.2, + seed=22, + drop_params=True, + ) # get a new response from the model where it can see the function response + print("second response\n", second_response) except litellm.InternalServerError as e: print(e) except litellm.RateLimitError as e: diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 8750d587dda..e608c2c73b6 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5523,14 +5523,22 @@ def test_neutralize_orphaned_tool_blocks_non_text_result_marked_not_empty(): "role": "assistant", "content": None, "tool_calls": [ - {"id": "c1", "type": "function", - "function": {"name": "render", "arguments": "{}"}} + { + "id": "c1", + "type": "function", + "function": {"name": "render", "arguments": "{}"}, + } ], }, { "role": "tool", "tool_call_id": "c1", - "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}], + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,AAAA"}, + } + ], }, ] @@ -5538,7 +5546,9 @@ def test_neutralize_orphaned_tool_blocks_non_text_result_marked_not_empty(): messages, optional_params={} ) - rewritten = next(m for m in result if m.get("role") == "user" and m is not messages[0]) + rewritten = next( + m for m in result if m.get("role") == "user" and m is not messages[0] + ) text = rewritten["content"] assert text.strip() # never empty assert "non-text tool result omitted" in text @@ -5756,3 +5766,56 @@ def test_transform_request_with_tools_still_builds_toolconfig(monkeypatch): ) assert "toolConfig" in result + + +def test_transform_request_flag_off_restores_raise(monkeypatch): + """Opt-out: with bedrock_neutralize_orphaned_tool_blocks=False and + modify_params=False, the legacy UnsupportedParamsError contract is restored.""" + monkeypatch.setattr(litellm, "bedrock_neutralize_orphaned_tool_blocks", False) + monkeypatch.setattr(litellm, "modify_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="without `tools="): + config.transform_request( + model="us.anthropic.claude-opus-4-5-20251101-v1:0", + messages=_orphaned_tool_history_messages(), + optional_params={}, + litellm_params={}, + headers={}, + ) + + +def test_transform_request_flag_off_with_modify_params_restores_dummy_tool(monkeypatch): + """Opt-out: with the flag off and modify_params=True, the legacy dummy-tool + injection is restored (a toolConfig is produced, not neutralized text).""" + monkeypatch.setattr(litellm, "bedrock_neutralize_orphaned_tool_blocks", False) + monkeypatch.setattr(litellm, "modify_params", True) + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-opus-4-5-20251101-v1:0", + messages=_orphaned_tool_history_messages(), + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert "toolConfig" in result + assert "dummy_tool" in json.dumps(result) + + +def test_transform_request_flag_on_is_default(monkeypatch): + """Default-on: without touching the flag, neutralization is the behavior.""" + monkeypatch.setattr(litellm, "modify_params", False) + config = AmazonConverseConfig() + + assert litellm.bedrock_neutralize_orphaned_tool_blocks is True + result = config.transform_request( + model="us.anthropic.claude-opus-4-5-20251101-v1:0", + messages=_orphaned_tool_history_messages(), + optional_params={}, + litellm_params={}, + headers={}, + ) + + _assert_no_structured_tool_blocks(result) From 681e75bf307284c51a40ee25e538a1e5b8a3a30c Mon Sep 17 00:00:00 2001 From: Kent Date: Fri, 26 Jun 2026 13:21:33 +0800 Subject: [PATCH 005/187] fix(bedrock): drop env-var backing for neutralize flag to satisfy doc CI The documentation_test_env_keys check requires every os.getenv() key to be documented in litellm-docs (a separate repo), which would create a cross-repo merge-order dependency. The flag does not need an env var: make it a plain module attribute (litellm.bedrock_neutralize_orphaned_tool_blocks = True), still overridable in Python or via proxy litellm_settings (which setattrs litellm module attrs). Default unchanged (True). --- litellm/__init__.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 4b10130326e..d8166280c68 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -242,10 +242,7 @@ telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False)) modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) -bedrock_neutralize_orphaned_tool_blocks = ( - os.getenv("LITELLM_BEDROCK_NEUTRALIZE_ORPHANED_TOOL_BLOCKS", "true").lower() - == "true" -) +bedrock_neutralize_orphaned_tool_blocks: bool = True use_chat_completions_url_for_anthropic_messages: bool = bool( os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False) ) # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API From 518c4b07a72e989ed777f93c01ef296c77a8b567 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 18 Jul 2026 21:55:44 +0000 Subject: [PATCH 006/187] fix(azure_ai): route Responses API to native /openai/v1/responses for Foundry Models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 5 + litellm/llms/azure_ai/chat/transformation.py | 8 +- litellm/llms/azure_ai/common_utils.py | 26 +++ litellm/llms/azure_ai/responses/__init__.py | 0 .../llms/azure_ai/responses/transformation.py | 62 +++++++ litellm/utils.py | 8 + .../llms/azure_ai/responses/__init__.py | 0 .../test_azure_ai_responses_transformation.py | 173 ++++++++++++++++++ 9 files changed, 279 insertions(+), 6 deletions(-) create mode 100644 litellm/llms/azure_ai/responses/__init__.py create mode 100644 litellm/llms/azure_ai/responses/transformation.py create mode 100644 tests/test_litellm/llms/azure_ai/responses/__init__.py create mode 100644 tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 2f6643c644c..6daf18e43ad 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1742,6 +1742,9 @@ if TYPE_CHECKING: from .llms.azure.responses.o_series_transformation import ( AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig, ) + from .llms.azure_ai.responses.transformation import ( + AzureAIResponsesAPIConfig as AzureAIResponsesAPIConfig, + ) from .llms.xai.responses.transformation import ( XAIResponsesAPIConfig as XAIResponsesAPIConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 488331e3895..744de95cab6 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -231,6 +231,7 @@ LLM_CONFIG_NAMES = ( "OpenAIResponsesAPIConfig", "AzureOpenAIResponsesAPIConfig", "AzureOpenAIOSeriesResponsesAPIConfig", + "AzureAIResponsesAPIConfig", "XAIResponsesAPIConfig", "LiteLLMProxyResponsesAPIConfig", "HostedVLLMResponsesAPIConfig", @@ -935,6 +936,10 @@ _LLM_CONFIGS_IMPORT_MAP = { ".llms.azure.responses.o_series_transformation", "AzureOpenAIOSeriesResponsesAPIConfig", ), + "AzureAIResponsesAPIConfig": ( + ".llms.azure_ai.responses.transformation", + "AzureAIResponsesAPIConfig", + ), "XAIResponsesAPIConfig": ( ".llms.xai.responses.transformation", "XAIResponsesAPIConfig", diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 27a98347087..27e4f405dfd 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -1,7 +1,6 @@ import enum import re from typing import Any, List, Optional, Tuple, cast -from urllib.parse import urlparse import httpx from httpx import Response @@ -12,6 +11,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( _audio_or_image_in_message_content, convert_content_list_to_str, ) +from litellm.llms.azure_ai.common_utils import azure_ai_use_api_key_header from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error @@ -85,11 +85,7 @@ class AzureAIStudioConfig(OpenAIConfig): """ Returns True if the request should use `api-key` header for authentication. """ - parsed_url = urlparse(api_base) - host = parsed_url.hostname - if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")): - return True - return False + return azure_ai_use_api_key_header(api_base) def get_complete_url( self, diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 9965aa693c3..26021367440 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,4 +1,5 @@ from typing import List, Literal, Optional +from urllib.parse import urlparse import litellm from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter @@ -6,6 +7,31 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues +def azure_ai_use_api_key_header(api_base: str) -> bool: + """Whether Azure AI auth should use the `api-key` header instead of a Bearer token. + + Foundry and Azure OpenAI hosts authenticate key-based requests with the + `api-key` header; serverless/other endpoints expect `Authorization: Bearer`. + """ + host = urlparse(api_base).hostname + return bool(host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com"))) + + +def azure_ai_supports_native_responses(model: str | None) -> bool: + """Whether an Azure AI model should use the native Responses API rather than the chat bridge. + + Foundry Models expose an OpenAI-compatible Responses endpoint at + `/openai/v1/responses`. Claude deployments speak the Anthropic + Messages API and the model-router/agents routes have their own surfaces, so + those keep the chat-completions bridge. + """ + if not model: + return False + if "claude" in model.lower(): + return False + return AzureFoundryModelInfo.get_azure_ai_route(model) == "default" + + class AzureFoundryModelInfo(BaseLLMModelInfo): """Model info for Azure AI / Azure Foundry models.""" diff --git a/litellm/llms/azure_ai/responses/__init__.py b/litellm/llms/azure_ai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/azure_ai/responses/transformation.py b/litellm/llms/azure_ai/responses/transformation.py new file mode 100644 index 00000000000..43fb9fc93d0 --- /dev/null +++ b/litellm/llms/azure_ai/responses/transformation.py @@ -0,0 +1,62 @@ +import httpx + +from litellm.llms.azure.common_utils import BaseAzureLLM +from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + azure_ai_use_api_key_header, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import _add_path_to_api_base + + +class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig): + """Native Responses API config for Azure AI Foundry Models. + + Foundry Models such as the GPT-5 family expose an OpenAI-compatible Responses + endpoint at `/openai/v1/responses`. Routing here (instead of the + chat-completions bridge) keeps `reasoning_effort` alongside function tools, + which Azure rejects on `/chat/completions`. + """ + + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.AZURE_AI + + def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict: + litellm_params = litellm_params or GenericLiteLLMParams() + api_key = AzureFoundryModelInfo.get_api_key(litellm_params.api_key) + api_base = AzureFoundryModelInfo.get_api_base(litellm_params.api_base) + + if api_key: + if api_base and azure_ai_use_api_key_header(api_base): + headers["api-key"] = api_key + else: + headers["Authorization"] = f"Bearer {api_key}" + else: + headers = BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params) + + headers.setdefault("Content-Type", "application/json") + return headers + + def get_complete_url( + self, + api_base: str | None, + litellm_params: dict, + ) -> str: + api_base = AzureFoundryModelInfo.get_api_base(api_base) + if api_base is None: + raise ValueError( + "api_base is required for Azure AI Foundry Responses API. " + "Set the api_base parameter or the AZURE_AI_API_BASE environment variable." + ) + + original_url = httpx.URL(api_base) + query_params = dict(original_url.params) + api_version = litellm_params.get("api_version") + if "api-version" not in query_params and isinstance(api_version, str): + query_params["api-version"] = api_version + + new_url = _add_path_to_api_base(api_base=api_base, ending_path="/openai/v1/responses") + return str(httpx.URL(new_url).copy_with(params=query_params)) diff --git a/litellm/utils.py b/litellm/utils.py index 174bed09396..34a083786d9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8221,6 +8221,14 @@ class ProviderConfigManager: return litellm.AzureOpenAIOSeriesResponsesAPIConfig() else: return litellm.AzureOpenAIResponsesAPIConfig() + elif litellm.LlmProviders.AZURE_AI == provider: + from litellm.llms.azure_ai.common_utils import ( + azure_ai_supports_native_responses, + ) + + if azure_ai_supports_native_responses(model): + return litellm.AzureAIResponsesAPIConfig() + return None elif litellm.LlmProviders.XAI == provider: return litellm.XAIResponsesAPIConfig() elif litellm.LlmProviders.GITHUB_COPILOT == provider: diff --git a/tests/test_litellm/llms/azure_ai/responses/__init__.py b/tests/test_litellm/llms/azure_ai/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py new file mode 100644 index 00000000000..0a2f2a4ab8e --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -0,0 +1,173 @@ +""" +Regression tests for native Azure AI Foundry Responses API routing (LIT-4427). + +Before the fix, `azure_ai` had no native Responses config, so `litellm.responses()` +fell back to the chat-completions bridge and sent `reasoning_effort` + function tools +to `/chat/completions`, which Azure rejects for GPT-5 models. These tests assert the +request now goes to the native `/openai/v1/responses` endpoint in Responses shape. +""" + +import json +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +import litellm +from litellm.llms.azure_ai.responses.transformation import AzureAIResponsesAPIConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ProviderConfigManager + + +class MockResponse: + def __init__(self, json_data, status_code=200): + self._json_data = json_data + self.status_code = status_code + self.text = json.dumps(json_data) + self.headers = httpx.Headers({}) + + def json(self): + return self._json_data + + +def _minimal_responses_payload(model: str) -> dict: + return { + "id": "resp_123", + "object": "response", + "created_at": 1741369938, + "status": "completed", + "model": model, + "output": [], + "parallel_tool_calls": False, + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + "error": None, + "tool_choice": "auto", + "tools": [], + "metadata": None, + "temperature": None, + "top_p": None, + "max_output_tokens": None, + "previous_response_id": None, + "reasoning": None, + "truncation": None, + "instructions": None, + "incomplete_details": None, + "user": None, + } + + +@pytest.mark.parametrize( + "model", + ["gpt-5.6-luna-20260710154139", "gpt-5.5-20260504143601", "DeepSeek-R1-0528"], +) +def test_azure_ai_resolves_native_responses_config(model): + config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model=model) + assert isinstance(config, AzureAIResponsesAPIConfig) + + +@pytest.mark.parametrize("model", ["claude-3-5-sonnet", "model_router/gpt-5", "agents/my-agent"]) +def test_azure_ai_non_responses_models_keep_bridge(model): + """Claude / model-router / agents routes have their own surfaces, so they must + keep returning None (chat-completions bridge).""" + config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model=model) + assert config is None + + +@pytest.mark.parametrize( + "api_base,expected", + [ + ( + "https://res.services.ai.azure.com/api/projects/proj", + "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses", + ), + ( + "https://res.services.ai.azure.com/api/projects/proj/", + "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses", + ), + ( + "https://res.services.ai.azure.com", + "https://res.services.ai.azure.com/openai/v1/responses", + ), + ( + "https://res.openai.azure.com", + "https://res.openai.azure.com/openai/v1/responses", + ), + ( + "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses", + "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses", + ), + ], +) +def test_get_complete_url(api_base, expected): + config = AzureAIResponsesAPIConfig() + assert config.get_complete_url(api_base=api_base, litellm_params={}) == expected + + +def test_validate_environment_api_key_header_for_foundry_host(): + config = AzureAIResponsesAPIConfig() + headers = config.validate_environment( + headers={}, + model="gpt-5.6-luna", + litellm_params=GenericLiteLLMParams( + api_key="secret", api_base="https://res.services.ai.azure.com/api/projects/proj" + ), + ) + assert headers["api-key"] == "secret" + assert "Authorization" not in headers + + +def test_validate_environment_bearer_for_serverless_host(): + config = AzureAIResponsesAPIConfig() + headers = config.validate_environment( + headers={}, + model="gpt-5.6-luna", + litellm_params=GenericLiteLLMParams( + api_key="secret", api_base="https://endpoint.eastus.models.ai.azure.com" + ), + ) + assert headers["Authorization"] == "Bearer secret" + assert "api-key" not in headers + + +@pytest.mark.asyncio +async def test_aresponses_routes_to_native_endpoint_with_reasoning_and_tools(): + """Core LIT-4427 regression: reasoning_effort + function tools must be sent to the + native /openai/v1/responses endpoint in Responses shape, not bridged to /chat/completions.""" + tools = [ + { + "type": "function", + "name": "get_weather", + "description": "Get weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + ] + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_post: + mock_post.return_value = MockResponse(_minimal_responses_payload("gpt-5.6-luna"), 200) + + await litellm.aresponses( + model="azure_ai/gpt-5.6-luna-20260710154139", + input="What is the weather in SF?", + reasoning_effort="high", + tools=tools, + api_base="https://res.services.ai.azure.com/api/projects/proj", + api_key="fake-key", + ) + + mock_post.assert_called_once() + url = str(mock_post.call_args.kwargs["url"]) + body = mock_post.call_args.kwargs["json"] + + assert url == "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses" + assert "/chat/completions" not in url + assert "input" in body + assert "messages" not in body + assert body["reasoning"] == {"effort": "high"} + assert body["tools"] == tools From 34d32c04e5cdf6a7ce88e8b3f359fa1b4d891d73 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 18 Jul 2026 22:06:58 +0000 Subject: [PATCH 007/187] test(azure_ai): drop responses test __init__ to fix package name collision; cover api-version, missing api_base, AD fallback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/azure_ai/responses/__init__.py | 0 .../test_azure_ai_responses_transformation.py | 35 +++++++++++++++++++ 2 files changed, 35 insertions(+) delete mode 100644 tests/test_litellm/llms/azure_ai/responses/__init__.py diff --git a/tests/test_litellm/llms/azure_ai/responses/__init__.py b/tests/test_litellm/llms/azure_ai/responses/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py index 0a2f2a4ab8e..b853b114645 100644 --- a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py +++ b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -103,6 +103,25 @@ def test_get_complete_url(api_base, expected): assert config.get_complete_url(api_base=api_base, litellm_params={}) == expected +def test_get_complete_url_adds_api_version_from_params(): + config = AzureAIResponsesAPIConfig() + url = config.get_complete_url( + api_base="https://res.services.ai.azure.com/api/projects/proj", + litellm_params={"api_version": "2025-04-01-preview"}, + ) + assert url == ( + "https://res.services.ai.azure.com/api/projects/proj/openai/v1/responses?api-version=2025-04-01-preview" + ) + + +def test_get_complete_url_raises_without_api_base(monkeypatch): + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("AZURE_AI_API_BASE", raising=False) + config = AzureAIResponsesAPIConfig() + with pytest.raises(ValueError): + config.get_complete_url(api_base=None, litellm_params={}) + + def test_validate_environment_api_key_header_for_foundry_host(): config = AzureAIResponsesAPIConfig() headers = config.validate_environment( @@ -129,6 +148,22 @@ def test_validate_environment_bearer_for_serverless_host(): assert "api-key" not in headers +def test_validate_environment_falls_back_to_base_azure_env_without_key(monkeypatch): + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False) + monkeypatch.delenv("AZURE_API_KEY", raising=False) + config = AzureAIResponsesAPIConfig() + headers = config.validate_environment( + headers={}, + model="gpt-5.6-luna", + litellm_params=GenericLiteLLMParams(api_base="https://res.services.ai.azure.com/api/projects/proj"), + ) + assert headers["Content-Type"] == "application/json" + assert "api-key" not in headers + + @pytest.mark.asyncio async def test_aresponses_routes_to_native_endpoint_with_reasoning_and_tools(): """Core LIT-4427 regression: reasoning_effort + function tools must be sent to the From 538a86885d11323a2883088947362d6b50f2e8fd Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 18 Jul 2026 22:17:03 +0000 Subject: [PATCH 008/187] fix(azure_ai): return native Responses config for management ops (model=None) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 2 +- .../responses/test_azure_ai_responses_transformation.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 34a083786d9..20d5993e14c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8226,7 +8226,7 @@ class ProviderConfigManager: azure_ai_supports_native_responses, ) - if azure_ai_supports_native_responses(model): + if model is None or azure_ai_supports_native_responses(model): return litellm.AzureAIResponsesAPIConfig() return None elif litellm.LlmProviders.XAI == provider: diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py index b853b114645..1d59773728a 100644 --- a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py +++ b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -65,6 +65,14 @@ def test_azure_ai_resolves_native_responses_config(model): assert isinstance(config, AzureAIResponsesAPIConfig) +def test_azure_ai_resolves_native_config_for_management_ops(): + """Management ops (delete/get/cancel/list) call the lookup with model=None; it must + still return the native config so those operations can build the right URL after a + native create succeeds.""" + config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model=None) + assert isinstance(config, AzureAIResponsesAPIConfig) + + @pytest.mark.parametrize("model", ["claude-3-5-sonnet", "model_router/gpt-5", "agents/my-agent"]) def test_azure_ai_non_responses_models_keep_bridge(model): """Claude / model-router / agents routes have their own surfaces, so they must From 1618617ea9f2b758f7ff95b8eb226579f739e522 Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 31 Jul 2026 21:47:22 +0000 Subject: [PATCH 009/187] feat(ui): accept ssh clone urls when registering a skill Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../skills/_components/add_plugin_form.tsx | 2 +- .../claude_code_plugins/helpers.test.ts | 43 +++++++++++++++ .../components/claude_code_plugins/helpers.ts | 52 +++++++++++++++++-- 3 files changed, 93 insertions(+), 4 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx index 04b8c88ae8d..226f9d19837 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx @@ -167,7 +167,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT label="Repository URL" name="skillUrl" rules={[{ required: true, message: "Please enter a repository URL" }]} - tooltip="Paste an HTTPS git repository URL from GitHub, GitLab, Bitbucket, or a self-hosted host. E.g. github.com/org/repo, gitlab.com/org/repo, or github.com/org/repo/tree/main/my-skill" + tooltip="Paste an HTTPS git repository URL from GitHub, GitLab, Bitbucket, or a self-hosted host, e.g. github.com/org/repo or github.com/org/repo/tree/main/my-skill. For a private repository use its SSH clone URL (git@ghe.example.com:org/repo.git) so Claude Code clones it with your own SSH key" > { it("returns null when no repo or url", () => { expect(getSourceLink({ source: "github" })).toBeNull(); }); + + it("returns null for an ssh clone url, which is not browsable", () => { + expect(getSourceLink({ source: "url", url: "git@ghe.example.com:org/repo.git" })).toBeNull(); + expect(getSourceLink({ source: "url", url: "ssh://git@ghe.example.com/org/repo.git" })).toBeNull(); + }); }); describe("getCategoryBadgeColor", () => { @@ -455,6 +460,44 @@ describe("parseSkillSource", () => { expect(parseSkillSource("gitlab.com/org/repo", "a//b")).toBeNull(); }); + it("keeps an scp-style ssh clone url so private hosts authenticate with the user's key", () => { + expect(parseSkillSource("git@ghe.example.com:org/repo.git")?.parsed).toEqual({ + source: "url", + url: "git@ghe.example.com:org/repo.git", + }); + expect(parseSkillSource("git@ghe.example.com:org/repo")?.parsed).toEqual({ + source: "url", + url: "git@ghe.example.com:org/repo.git", + }); + expect(parseSkillSource("git@ghe.example.com:org/repo.git")?.suggestedName).toBe("repo"); + }); + + it("normalizes an ssh:// clone url and keeps a custom port", () => { + expect(parseSkillSource("ssh://git@ghe.example.com/org/repo")?.parsed).toEqual({ + source: "url", + url: "ssh://git@ghe.example.com/org/repo.git", + }); + expect(parseSkillSource("ssh://git@ghe.example.com:2222/org/nested/repo.git")?.parsed).toEqual({ + source: "url", + url: "ssh://git@ghe.example.com:2222/org/nested/repo.git", + }); + }); + + it("combines an ssh clone url with an explicit subfolder", () => { + expect(parseSkillSource("git@ghe.example.com:org/repo.git", "plugins/my-skill")?.parsed).toEqual({ + source: "git-subdir", + url: "git@ghe.example.com:org/repo.git", + path: "plugins/my-skill", + }); + expect(parseSkillSource("git@ghe.example.com:org/repo.git", "../etc")).toBeNull(); + }); + + it("rejects ssh-looking input without a host or repo path", () => { + expect(parseSkillSource("git@ghe.example.com:repo.git")).toBeNull(); + expect(parseSkillSource("git@localhost:org/repo.git")).toBeNull(); + expect(parseSkillSource("git@:org/repo.git")).toBeNull(); + }); + it("returns null for empty and garbage input", () => { expect(parseSkillSource("")).toBeNull(); expect(parseSkillSource(" ")).toBeNull(); diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts index a4e70f78af1..caf59c71b76 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts @@ -30,6 +30,10 @@ const IPV4_HOST_REGEX = /^\d{1,3}(\.\d{1,3}){3}$/; const GITHUB_ORG_REGEX = /^[A-Za-z0-9-]+$/; const GITHUB_REPO_REGEX = /^[A-Za-z0-9._-]+$/; +const SSH_SCP_REGEX = /^([a-z0-9._-]+)@([a-z0-9.-]+\.[a-z]{2,}):([a-z0-9._-]+(?:\/[a-z0-9._-]+)+?)(?:\.git)?\/?$/i; +const SSH_URL_REGEX = + /^ssh:\/\/([a-z0-9._-]+)@([a-z0-9.-]+\.[a-z]{2,})(:\d+)?\/([a-z0-9._-]+(?:\/[a-z0-9._-]+)+?)(?:\.git)?\/?$/i; + const buildRepoUrl = (url: URL): string => `${url.protocol}//${url.host}${url.pathname.replace(/\/+$/, "")}`; const pathSegments = (url: URL): string[] => url.pathname.split("/").filter((seg) => seg !== ""); @@ -160,12 +164,54 @@ const parseRawGitSource = (url: URL, subPath?: string): SkillSourcePreview | nul }; }; +const withGitSuffix = (path: string): string => `${path.replace(/\.git$/i, "")}.git`; + +const parseSshRepoUrl = (raw: string): string | null => { + const trimmed = raw.trim(); + const sshUrl = SSH_URL_REGEX.exec(trimmed); + if (sshUrl) { + const [, user, host, port, path] = sshUrl; + return `ssh://${user}@${host}${port ?? ""}/${withGitSuffix(path)}`; + } + const scp = SSH_SCP_REGEX.exec(trimmed); + if (scp) { + const [, user, host, path] = scp; + return `${user}@${host}:${withGitSuffix(path)}`; + } + return null; +}; + +const parseSshSource = (cloneUrl: string, subPath?: string): SkillSourcePreview | null => { + const repoName = lastSegment(cloneUrl.replace(/\.git$/, "").replace(/^[^:]*:/, "")); + const normalized = normalizeSubPath(subPath ?? ""); + if (normalized !== "") { + if (!SUBDIR_PATH_REGEX.test(normalized)) { + return null; + } + return { + parsed: { source: "git-subdir", url: cloneUrl, path: normalized }, + label: `SSH subdir — ${cloneUrl} @ ${normalized}`, + suggestedName: toKebabCase(lastSegment(normalized)), + }; + } + return { + parsed: { source: "url", url: cloneUrl }, + label: `SSH repo — ${cloneUrl}`, + suggestedName: toKebabCase(repoName), + }; +}; + /** * Parse any git-accessible repository URL into a registerable skill source. - * GitHub URLs keep their `github`/`git-subdir` shorthand; every other host is - * treated as a raw repo URL, with an optional subfolder turning it into git-subdir. + * GitHub https URLs keep their `github`/`git-subdir` shorthand; ssh clone URLs stay ssh so a + * private host authenticates with the user's own key; every other host is treated as a raw repo + * URL, with an optional subfolder turning it into git-subdir. */ export const parseSkillSource = (rawUrl: string, subPath?: string): SkillSourcePreview | null => { + const sshCloneUrl = parseSshRepoUrl(rawUrl); + if (sshCloneUrl) { + return parseSshSource(sshCloneUrl, subPath); + } const url = parseRepoUrl(rawUrl); if (!url) { return null; @@ -268,7 +314,7 @@ export const getSourceLink = (source: PluginSource): string | null => { return `https://github.com/${source.repo}`; } if ((source.source === "url" || source.source === "git-subdir") && source.url) { - return source.url; + return source.url.startsWith("https://") ? source.url : null; } return null; }; From 66e1ea6090e5e97a5fa6c17aba376a1d57bc6e66 Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 31 Jul 2026 21:58:27 +0000 Subject: [PATCH 010/187] fix(ui): render non-browsable skill sources as text on the detail page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../claude_code_plugins/skill_detail.test.tsx | 42 +++++++++++++++++++ .../claude_code_plugins/skill_detail.tsx | 21 +++++++--- 2 files changed, 57 insertions(+), 6 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx new file mode 100644 index 00000000000..98ae4a44ae6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx @@ -0,0 +1,42 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import { Plugin } from "./types"; + +import SkillDetail from "./skill_detail"; + +const buildSkill = (source: Plugin["source"]): Plugin => ({ + id: "plugin-id", + name: "my-skill", + source, + enabled: true, +}); + +describe("SkillDetail source", () => { + it("links a github source to the repository", () => { + render(); + expect(screen.getByRole("link", { name: /github.com\/org\/repo/ })).toHaveAttribute( + "href", + "https://github.com/org/repo", + ); + }); + + it("renders an ssh clone url as plain text instead of an unusable link", () => { + render( + , + ); + expect(screen.getByText("git@ghe.example.com:org/repo.git")).toBeInTheDocument(); + expect(screen.queryByRole("link", { name: /ghe.example.com/ })).not.toBeInTheDocument(); + }); + + it("renders an ssh git-subdir source as plain text without a tree path", () => { + render( + , + ); + expect(screen.getByText("git@ghe.example.com:org/repo.git")).toBeInTheDocument(); + expect(screen.queryByRole("link", { name: /ghe.example.com/ })).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx index fe001641135..0f25b8ad515 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx @@ -1,6 +1,6 @@ import React, { useState } from "react"; import { ArrowLeftOutlined, CopyOutlined, CheckOutlined, LinkOutlined } from "@ant-design/icons"; -import { buildMarketplaceSettingsSnippet, formatInstallCommand } from "./helpers"; +import { buildMarketplaceSettingsSnippet, formatInstallCommand, getSourceLink } from "./helpers"; import { Plugin } from "./types"; interface SkillDetailProps { @@ -21,13 +21,15 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { setTimeout(() => setCopiedKey(null), 2000); }; + const sourceLink = getSourceLink(skill.source); const sourceUrl = (() => { const src = skill.source; - if (src.source === "github" && src.repo) return `https://github.com/${src.repo}`; - if (src.source === "git-subdir" && src.url) return src.path ? `${src.url}/tree/main/${src.path}` : src.url; - if (src.source === "url" && src.url) return src.url; - return null; + if (sourceLink && src.source === "git-subdir" && src.path) { + return `${sourceLink}/tree/main/${src.path}`; + } + return sourceLink; })(); + const sourceText = skill.source.url ?? sourceUrl; const installCommand = formatInstallCommand(skill); @@ -146,7 +148,7 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { - {sourceUrl && ( + {sourceUrl ? ( + ) : ( + sourceText && ( +
+
Source
+
{sourceText}
+
+ ) )} {skill.keywords && skill.keywords.length > 0 && ( From 7d5c5d8873e55cb47b82a351046f2af31be06cf3 Mon Sep 17 00:00:00 2001 From: milan Date: Fri, 31 Jul 2026 22:03:38 +0000 Subject: [PATCH 011/187] fix(ui): show the subfolder path for non-browsable skill sources Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/components/claude_code_plugins/skill_detail.test.tsx | 2 +- .../src/components/claude_code_plugins/skill_detail.tsx | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx index 98ae4a44ae6..2f600397a4b 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx @@ -36,7 +36,7 @@ describe("SkillDetail source", () => { onBack={vi.fn()} />, ); - expect(screen.getByText("git@ghe.example.com:org/repo.git")).toBeInTheDocument(); + expect(screen.getByText("git@ghe.example.com:org/repo.git @ plugins/x")).toBeInTheDocument(); expect(screen.queryByRole("link", { name: /ghe.example.com/ })).not.toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx index 0f25b8ad515..ad98fdf232c 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx @@ -29,7 +29,9 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { } return sourceLink; })(); - const sourceText = skill.source.url ?? sourceUrl; + const sourceText = skill.source.url + ? `${skill.source.url}${skill.source.path ? ` @ ${skill.source.path}` : ""}` + : sourceUrl; const installCommand = formatInstallCommand(skill); From 62d8258868b745bfab8f7521ba498ad2a1564b06 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:27:50 +0000 Subject: [PATCH 012/187] fix(dashscope): forward reasoning_effort to the provider Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/dashscope/chat/transformation.py | 6 ++++++ .../test_dashscope_chat_transformation.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index 5ab7fbf3658..977bb38f59a 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -12,6 +12,12 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig class DashScopeChatConfig(OpenAIGPTConfig): + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class contract returns a list + return [ # mutable-ok: base class contract returns a list + *super().get_supported_openai_params(model=model), + "reasoning_effort", + ] + def remove_cache_control_flag_from_messages_and_tools( self, model: str, diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py b/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py index 8dbc197d4b5..e3bbf2abc48 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_chat_transformation.py @@ -169,6 +169,21 @@ class TestDashScopeConfig: assert transformed_messages[0].get("cache_control") == {"type": "ephemeral"} + @pytest.mark.parametrize("reasoning_effort", ["none", "minimal", "low", "high"]) + def test_dashscope_forwards_reasoning_effort(self, reasoning_effort: str): + """DashScope supports reasoning_effort, so it must reach the provider instead of being dropped.""" + assert "reasoning_effort" in DashScopeChatConfig().get_supported_openai_params( + model="qwen3.7-plus" + ) + + optional_params = litellm.get_optional_params( + model="qwen3.7-plus", + custom_llm_provider="dashscope", + reasoning_effort=reasoning_effort, + ) + + assert optional_params["reasoning_effort"] == reasoning_effort + def test_dashscope_preserves_cache_control_in_tools(self): """DashScope should NOT strip cache_control from tools.""" config = DashScopeChatConfig() From 2b6c30c2d93199586fd696fab7cedc0c5f769b16 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 20:42:38 +0000 Subject: [PATCH 013/187] refactor(dashscope): tighten supported params return type Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/dashscope/chat/transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index 977bb38f59a..04d530ea89b 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -12,7 +12,7 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig class DashScopeChatConfig(OpenAIGPTConfig): - def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class contract returns a list + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: base class contract returns a list return [ # mutable-ok: base class contract returns a list *super().get_supported_openai_params(model=model), "reasoning_effort", From dff08dcb55b35ff9346445b0ba3a733692e249f1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:16:45 +0000 Subject: [PATCH 014/187] fix(proxy): retry rate-limit fallbacks from a pristine request snapshot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 38 ++- .../proxy/test_common_request_processing.py | 240 ++++++++++++++++++ 2 files changed, 264 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e3a2b892721..255bf3ebda2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -32,7 +32,11 @@ from litellm.constants import ( UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error +from litellm.litellm_core_utils.core_helpers import ( + get_or_create_metadata_bucket, + independent_snapshot, + is_expected_client_error, +) from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, @@ -2034,6 +2038,21 @@ class ProxyBaseLLMRequestProcessing: ) -> tuple[dict, LiteLLMLoggingObj]: from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + original_model: Final = self.data.get("model") + fallback_models: Final = ( + self._resolve_fallback_models( + model=original_model, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) + if original_model + and isinstance(original_model, str) + and llm_router + and not self.data.get("disable_fallbacks") + else None + ) + pristine: Final = independent_snapshot(self.data) if fallback_models else None + try: return await self.common_processing_pre_call_logic( request=request, @@ -2052,16 +2071,7 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, ) except ProxyRateLimitError as original_exc: - original_model: Final = self.data.get("model") - if not original_model or not llm_router or self.data.get("disable_fallbacks"): - raise - - fallback_models: Final = self._resolve_fallback_models( - model=original_model, - llm_router=llm_router, - user_api_key_dict=user_api_key_dict, - ) - if not fallback_models: + if not fallback_models or pristine is None: raise verbose_proxy_logger.info( @@ -2074,7 +2084,7 @@ class ProxyBaseLLMRequestProcessing: for fallback_model in fallback_models: if fallback_model == original_model: continue - self.data["model"] = fallback_model + self.data = {**independent_snapshot(pristine), "model": fallback_model} try: return await self.common_processing_pre_call_logic( request=request, @@ -2095,10 +2105,10 @@ class ProxyBaseLLMRequestProcessing: except ProxyRateLimitError: continue except BaseException: - self.data["model"] = original_model + self.data = pristine raise - self.data["model"] = original_model + self.data = pristine raise original_exc def _resolve_fallback_models( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 50b26577e5c..89e0799ba16 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6235,6 +6235,246 @@ class TestPreCallWithFallbacksOnLocalRateLimit: call_type="acompletion", ) + @pytest.mark.asyncio + async def test_fallback_retries_from_pristine_request_data(self): + import threading + + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model = "gpt-4" + fallback_model = "gpt-3.5-turbo" + + processor = ProxyBaseLLMRequestProcessing( + data={ + "model": primary_model, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"tags": ["a"]}, + } + ) + + metadata_at_entry = [] + + async def mock_pre_call_logic(**kwargs): + copy.deepcopy(processor.data["metadata"]) + metadata_at_entry.append(dict(processor.data["metadata"])) + processor.data["metadata"]["litellm_parent_otel_span"] = threading.RLock() + processor.data["litellm_logging_obj"] = object() + if processor.data.get("model") == primary_model: + raise ProxyRateLimitError( + detail="TPM limit exceeded for gpt-4", + headers={"retry-after": "30"}, + ) + return processor.data, MagicMock() + + mock_router = MagicMock() + mock_router.fallbacks = [{primary_model: [fallback_model]}] + + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=mock_pre_call_logic, + ): + data, logging_obj = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=primary_model, + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data["model"] == fallback_model + assert metadata_at_entry[1] == {"tags": ["a"]} + + @pytest.mark.asyncio + async def test_exhausted_fallbacks_restore_pristine_request_data(self): + import threading + + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model = "gpt-4" + original_data = { + "model": primary_model, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"tags": ["a"]}, + } + processor = ProxyBaseLLMRequestProcessing(data=copy.deepcopy(original_data)) + + async def mock_pre_call_logic(**kwargs): + processor.data["metadata"]["litellm_parent_otel_span"] = threading.RLock() + processor.data["litellm_logging_obj"] = object() + raise ProxyRateLimitError( + detail=f"TPM limit exceeded for {processor.data.get('model')}", + headers={"retry-after": "30"}, + ) + + mock_router = MagicMock() + mock_router.fallbacks = [{primary_model: ["gpt-3.5-turbo"]}] + + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=mock_pre_call_logic, + ): + with pytest.raises(ProxyRateLimitError, match="gpt-4"): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=primary_model, + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data == original_data + + @pytest.mark.asyncio + async def test_real_add_litellm_data_to_request_rerun_with_otel_span_falls_back(self): + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + from litellm.proxy.proxy_server import ProxyConfig + + trace.set_tracer_provider(TracerProvider()) + + primary_model = "gpt-4" + fallback_model = "gpt-3.5-turbo" + + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = "/v1/chat/completions" + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + user_api_key_dict = UserAPIKeyAuth( + parent_otel_span=trace.get_tracer("x").start_span("s"), + api_key="hashed-key", + user_id="u1", + team_id="t1", + metadata={}, + team_metadata={}, + team_member_tpm_limit=1000, + ) + + processor = ProxyBaseLLMRequestProcessing( + data={ + "model": primary_model, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"tags": ["a"]}, + } + ) + + async def real_add_litellm_data_pre_call(**kwargs): + await add_litellm_data_to_request( + data=processor.data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=ProxyConfig(), + general_settings={}, + version="test", + ) + if processor.data.get("model") == primary_model: + raise ProxyRateLimitError( + detail="TPM limit exceeded for gpt-4", + headers={"retry-after": "30"}, + ) + return processor.data, MagicMock() + + mock_router = MagicMock() + mock_router.fallbacks = [{primary_model: [fallback_model]}] + + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=real_add_litellm_data_pre_call, + ): + data, logging_obj = await processor._pre_call_with_fallbacks( + request=request_mock, + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=user_api_key_dict, + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=primary_model, + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data["model"] == fallback_model + + @pytest.mark.asyncio + async def test_no_fallbacks_skips_snapshot(self): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + raise ProxyRateLimitError( + detail="TPM limit exceeded", + headers={"retry-after": "30"}, + ) + + mock_router = MagicMock() + mock_router.fallbacks = None + + with patch("litellm.proxy.common_request_processing.independent_snapshot") as snapshot_mock: + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=mock_pre_call_logic, + ): + with pytest.raises(ProxyRateLimitError): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + snapshot_mock.assert_not_called() + class _RecordingSuccessLogger(CustomLogger): def __init__(self): From c064e576ee31ed05b067462412c6c234d364c8fd Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:29:20 +0000 Subject: [PATCH 015/187] fix(proxy): tolerate missing router_settings and non-list fallbacks in fallback resolution Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 4 ++-- tests/test_litellm/proxy/test_common_request_processing.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 255bf3ebda2..9927320d795 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2121,14 +2121,14 @@ class ProxyBaseLLMRequestProcessing: fallbacks = None - key_router_settings: Final = user_api_key_dict.router_settings + key_router_settings: Final = getattr(user_api_key_dict, "router_settings", None) if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings: fallbacks = key_router_settings["fallbacks"] if fallbacks is None: fallbacks = llm_router.fallbacks - if not fallbacks: + if not isinstance(fallbacks, list) or not fallbacks: return None fallback_model_group, generic_fallback_idx = get_fallback_model_group( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 89e0799ba16..4beb73b3a81 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6449,7 +6449,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router = MagicMock() mock_router.fallbacks = None - with patch("litellm.proxy.common_request_processing.independent_snapshot") as snapshot_mock: + with patch( # test-quality-ok: spying the snapshot seam is the only observable check that the no-fallback path skips it + "litellm.proxy.common_request_processing.independent_snapshot" + ) as snapshot_mock: with patch.object( processor, "common_processing_pre_call_logic", From 431dcce6a72ddcb6ba73c56848e1b5e0477b7699 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:46:44 +0000 Subject: [PATCH 016/187] test(proxy): give pre-call mocks real router_settings and fallbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 4 ++-- .../test_response_polling_pre_call_checks.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9927320d795..255bf3ebda2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2121,14 +2121,14 @@ class ProxyBaseLLMRequestProcessing: fallbacks = None - key_router_settings: Final = getattr(user_api_key_dict, "router_settings", None) + key_router_settings: Final = user_api_key_dict.router_settings if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings: fallbacks = key_router_settings["fallbacks"] if fallbacks is None: fallbacks = llm_router.fallbacks - if not isinstance(fallbacks, list) or not fallbacks: + if not fallbacks: return None fallback_model_group, generic_fallback_idx = get_fallback_model_group( diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py index 459834d0fd2..38f087f51ca 100644 --- a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py +++ b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py @@ -48,10 +48,10 @@ class TestSkipPreCallLogic: await processor.base_process_llm_request( request=MagicMock(spec=Request), fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth, router_settings=None), route_type="aresponses", proxy_logging_obj=mock_proxy_logging, - llm_router=MagicMock(), + llm_router=MagicMock(fallbacks=None), general_settings={}, proxy_config=MagicMock(), skip_pre_call_logic=True, @@ -87,10 +87,10 @@ class TestSkipPreCallLogic: await processor.base_process_llm_request( request=MagicMock(spec=Request), fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth, router_settings=None), route_type="aresponses", proxy_logging_obj=mock_proxy_logging, - llm_router=MagicMock(), + llm_router=MagicMock(fallbacks=None), general_settings={}, proxy_config=MagicMock(), ) From 5493be743b1487e28c08868f5e8e920a3a8d2f7a Mon Sep 17 00:00:00 2001 From: Kent Date: Sat, 12 Sep 2026 09:33:07 +0800 Subject: [PATCH 017/187] refactor(bedrock): fold orphaned-tool handling into existing rebind to satisfy LIT011 budget --- litellm/llms/bedrock/chat/converse_transformation.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index d0d97ecb7f0..4c618b0ba83 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1793,10 +1793,10 @@ class AmazonConverseConfig(BaseConfig): ) -> RequestObject: messages, system_content_blocks = self._transform_system_message(messages, model=model) - messages = self._handle_orphaned_tool_blocks(messages, optional_params) - # Convert last user message to guarded_text if guardrailConfig is present - messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) + messages = self._convert_consecutive_user_messages_to_guarded_text( + self._handle_orphaned_tool_blocks(messages, optional_params), optional_params + ) ## TRANSFORMATION ## _data: Final[CommonRequestObject] = self._transform_request_helper( @@ -1856,10 +1856,10 @@ class AmazonConverseConfig(BaseConfig): ) -> RequestObject: messages, system_content_blocks = self._transform_system_message(messages, model=model) - messages = self._handle_orphaned_tool_blocks(messages, optional_params) - # Convert last user message to guarded_text if guardrailConfig is present - messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) + messages = self._convert_consecutive_user_messages_to_guarded_text( + self._handle_orphaned_tool_blocks(messages, optional_params), optional_params + ) _data: Final[CommonRequestObject] = self._transform_request_helper( model=model, From 04086be7955615a2498ab36e1f6539447be9ab73 Mon Sep 17 00:00:00 2001 From: Kent Date: Sat, 12 Sep 2026 09:51:01 +0800 Subject: [PATCH 018/187] fix(bedrock): guard neutralized tool results anywhere in history when guardrailConfig is set --- .../bedrock/chat/converse_transformation.py | 18 ++++--- .../chat/test_converse_transformation.py | 48 +++++++++++++++++++ 2 files changed, 60 insertions(+), 6 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 4c618b0ba83..e44dece3a8b 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -234,22 +234,28 @@ class AmazonConverseConfig(BaseConfig): rendered = convert_content_list_to_str(message).strip() return rendered or "" + guardrail_active: Final = "guardrailConfig" in optional_params + def _rewrite(message: AllMessageValues) -> AllMessageValues: role = message.get("role") tool_calls = message.get("tool_calls") if role == "assistant" and tool_calls: - base_text = convert_content_list_to_str(message) - call_texts = [_tool_call_text(call) for call in tool_calls] - text = "\n".join(filter(None, [base_text, *call_texts])) + base_text: Final = convert_content_list_to_str(message) + call_texts: Final = tuple(_tool_call_text(call) for call in tool_calls) + text: Final = "\n".join(part for part in (base_text, *call_texts) if part) return ChatCompletionAssistantMessage(role="assistant", content=text) if role in ("tool", "function"): tool_call_id = message.get("tool_call_id") name = message.get("name") label = f"tool result for {tool_call_id or name or 'unknown'}" - return ChatCompletionUserMessage( - role="user", - content=f"[{label}: {_result_text(message)}]", + result_text: Final = f"[{label}: {_result_text(message)}]" + # Tool results are externally controlled, so guard them wherever they + # land in history; _convert_consecutive_user_messages_to_guarded_text + # only covers the trailing user turn. + content: Final = ( + [{"type": "guarded_text", "text": result_text}] if guardrail_active else result_text ) + return ChatCompletionUserMessage(role="user", content=content) return message verbose_logger.warning( diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index ca550530369..21e8e94d871 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -6714,6 +6714,54 @@ def test_transform_request_neutralized_tool_output_is_guarded(monkeypatch): assert "secret tool output" in serialized +def test_transform_request_neutralized_tool_output_guarded_mid_history(monkeypatch): + """Regression: a neutralized tool result that is NOT the trailing turn (an + assistant reply and a later user turn follow it) must still be guardContent. + _convert_consecutive_user_messages_to_guarded_text only covers the trailing + user turn, so neutralize itself must guard untrusted tool output regardless + of position, else an attacker controlling the tool response bypasses the + guardrail (bot review).""" + monkeypatch.setattr(litellm, "modify_params", False) + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-opus-4-5-20251101-v1:0", + messages=[ + {"role": "user", "content": "look it up"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "c1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "IGNORE_PRIOR malware"}, + {"role": "assistant", "content": "Here is the summary."}, + {"role": "user", "content": "thanks"}, + ], + optional_params={ + "guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "1"} + }, + litellm_params={}, + headers={}, + ) + + _assert_no_structured_tool_blocks(result) + blocks = [block for message in result["messages"] for block in message["content"]] + guarded_texts = [ + block["guardContent"]["text"]["text"] for block in blocks if "guardContent" in block + ] + plain_texts = [block["text"] for block in blocks if "text" in block and "guardContent" not in block] + assert any("malware" in text for text in guarded_texts), "mid-history tool output must be guarded" + assert not any( + "malware" in text for text in plain_texts + ), "mid-history tool output must not reach the model as unguarded text" + + @pytest.mark.asyncio async def test_async_transform_request_no_tools_with_tool_history(monkeypatch): """Async is a separate request assembler; it must neutralize identically.""" From e250e661e36f217f338e0fe2e4e51415a45847d5 Mon Sep 17 00:00:00 2001 From: Kent Date: Sat, 12 Sep 2026 10:02:36 +0800 Subject: [PATCH 019/187] style(bedrock): apply ruff format to neutralize helper --- litellm/llms/bedrock/chat/converse_transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index e44dece3a8b..33a77e8d346 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -252,9 +252,7 @@ class AmazonConverseConfig(BaseConfig): # Tool results are externally controlled, so guard them wherever they # land in history; _convert_consecutive_user_messages_to_guarded_text # only covers the trailing user turn. - content: Final = ( - [{"type": "guarded_text", "text": result_text}] if guardrail_active else result_text - ) + content: Final = [{"type": "guarded_text", "text": result_text}] if guardrail_active else result_text return ChatCompletionUserMessage(role="user", content=content) return message From 07bed091194a4b49a4440357da3f106bd7798e5d Mon Sep 17 00:00:00 2001 From: Elif Naz Ozdamar <83784925+elifozdamar@users.noreply.github.com> Date: Sat, 12 Sep 2026 08:45:15 +0000 Subject: [PATCH 020/187] fix(proxy): release completed max-parallel slots promptly --- .../hooks/parallel_request_limiter_v3.py | 78 +++++------ .../hooks/test_parallel_request_limiter_v3.py | 121 ++++++++++++++++++ 2 files changed, 153 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c398abff099..55638ab9071 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -522,6 +522,7 @@ class RequestRateLimiterStash: owner_litellm_call_id: str | None = None rate_limit_response: RateLimitResponse | None = None parallel_slot: ParallelSlotAcquisition | None = None + parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False) reserved_tokens: int = 0 reserved_model: str | None = None reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) @@ -1609,6 +1610,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): statuses.append(self._gauge_status(gauge, in_flight + 1, "OK")) return RateLimitResponse(overall_code="OK", statuses=statuses) + async def _release_stashed_parallel_slot( + self, + stash: RequestRateLimiterStash | None, + parent_otel_span: Span | None, + ) -> None: + if stash is None: + return + async with stash.parallel_slot_release_lock: + acquisition: Final = stash.parallel_slot + if acquisition is None: + return + await self._release_parallel_request_slots(acquisition, parent_otel_span) + stash.parallel_slot = None # rebind-ok: marks this request's slot as released + async def _release_parallel_request_slots( self, acquisition: ParallelSlotAcquisition, @@ -3368,13 +3383,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.reservation_released = True - acquisition: Final = stash.parallel_slot - if acquisition is not None: - await self._release_parallel_request_slots( - acquisition=acquisition, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=io_response, descriptors=descriptors, @@ -3631,13 +3640,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if tpm_response["overall_code"] == "OVER_LIMIT": - acquisition: Final = stash.parallel_slot - if acquisition is not None: - await self._release_parallel_request_slots( - acquisition=acquisition, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=tpm_response, descriptors=descriptors, @@ -4450,13 +4453,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) - acquisition: Final = stash.parallel_slot if stash is not None else None - if stash is not None and acquisition is not None: - await self._release_parallel_request_slots( - acquisition=acquisition, - parent_otel_span=litellm_parent_otel_span, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span) pipeline_operations: Final = self._build_success_event_pipeline_operations( kwargs=kwargs, @@ -4576,13 +4573,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [] stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) - acquisition: Final = stash.parallel_slot if stash is not None else None - if stash is not None and acquisition is not None: - await self._release_parallel_request_slots( - acquisition=acquisition, - parent_otel_span=litellm_parent_otel_span, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span) # Skip the reservation refund if async_post_call_failure_hook # already released it (proxy-level rejection that also bubbles up @@ -4690,23 +4681,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): object's current max_parallel_requests configuration, which can change mid-request) decides whether there is anything to release. """ - stash: Final = get_request_stash() - if stash is None or stash.parallel_slot is None: - return - - await self._release_parallel_request_slots( - acquisition=stash.parallel_slot, - parent_otel_span=None, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(get_request_stash(), None) async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ - Post-call hook to update rate limit headers in the response. + Release completed-request slots and update rate limit headers in the response. """ try: - stash: Final = get_request_stash() - litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None + slot_stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(data)) + await self._release_stashed_parallel_slot(slot_stash, user_api_key_dict.parent_otel_span) + except Exception as e: + verbose_proxy_logger.exception("Error releasing parallel request slot in post-call hook: %s", e) + + try: + header_stash: Final = get_request_stash() + litellm_proxy_rate_limit_response: Final = ( + header_stash.rate_limit_response if header_stash is not None else None + ) if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response): additional_headers: Final = ensure_response_additional_headers(response) @@ -4774,12 +4765,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): stash: Final = get_request_stash() if stash is None: return - if stash.parallel_slot is not None: - await self._release_parallel_request_slots( - acquisition=stash.parallel_slot, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) if stash.batch_enqueued_reservation is not None: await self.batch_enqueued_token_store.refund( diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 10c0bb88a82..86f0d76e063 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7,6 +7,7 @@ import logging import os import sys import time +from collections.abc import Sequence from contextlib import contextmanager from datetime import datetime, timedelta from typing import Any, Dict, List, Optional @@ -32,6 +33,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( ) from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ( EmbeddingResponse, ModelResponse, @@ -4054,6 +4056,125 @@ async def _seed_max_parallel_requests_slots( ) +@pytest.mark.asyncio +async def test_completed_responses_post_call_releases_parallel_slot() -> None: + api_key = hash_token("sk-responses-post-call") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + data = { + "model": "gpt-4o-mini", + "input": "hello", + "litellm_call_id": "responses-owner", + } + parallel_key = f"{{api_key:{api_key}}}:max_parallel_requests" + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type="aresponses", + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 1 + + await handler.async_post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=ResponsesAPIResponse( + id="resp_parallel_slot", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[], + status="completed", + ), + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 0 + + await handler.async_log_success_event( + kwargs={"litellm_call_id": data["litellm_call_id"]}, + response_obj=None, + start_time=None, + end_time=None, + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 0 + + +@pytest.mark.asyncio +async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None: + from unittest.mock import AsyncMock + + api_key = hash_token("sk-concurrent-release") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) + call_id = "concurrent-release-owner" + parallel_key = f"{{api_key:{api_key}}}:max_parallel_requests" + release_started = asyncio.Event() + allow_redis_failure = asyncio.Event() + + async def failing_release( + keys: Sequence[str], args: Sequence[object] + ) -> list[int]: + release_started.set() + await allow_redis_failure.wait() + raise ConnectionError("redis unavailable") + + release_script = AsyncMock(side_effect=failing_release) + handler.parallel_release_script = release_script + await local_cache.async_set_cache(key=parallel_key, value=2, local_only=True) + stash = get_or_create_request_stash() + stash.owner_litellm_call_id = call_id + stash.parallel_slot = ParallelSlotAcquisition( + slot_id="slot-concurrent-release", + counter_keys=[parallel_key], + ) + data = {"litellm_call_id": call_id} + + post_call_task = asyncio.create_task( + handler.async_post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=ResponsesAPIResponse( + id="resp_concurrent_release", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[], + status="completed", + ), + ) + ) + await asyncio.wait_for(release_started.wait(), timeout=5) + logging_task = asyncio.create_task( + handler.async_log_success_event( + kwargs=data, + response_obj=None, + start_time=None, + end_time=None, + ) + ) + allow_redis_failure.set() + await asyncio.wait_for( + asyncio.gather(post_call_task, logging_task), + timeout=5, + ) + + assert release_script.await_count == 1 + assert await local_cache.async_get_cache(key=parallel_key) == 1 + assert stash.parallel_slot is None + + async def _build_seeded_limiter(): """Build a v3 limiter whose api-key slot registry already holds the pre-call slot.""" api_key = hash_token("sk-disconnect") From 682296ad68ce2afce1bc3cece84b70f88c86a8ac Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 05:00:32 -0700 Subject: [PATCH 021/187] test: verify database transactions and persisted accounting contracts --- .circleci/config.yml | 2 +- tests/e2e/coverage_registry/other.yaml | 4 + .../coverage_registry/quota_management.yaml | 6 + tests/integration/README.md | 6 +- tests/integration/_support/process.py | 92 ++++++++++ tests/integration/contracts.json | 28 +++ .../database/test_partition_transactions.py | 86 +++++++++ .../test_reader_writer_regeneration.py | 91 ++++++++++ .../database/test_transaction_atomicity.py | 67 +++++++ .../pricing/test_price_precedence.py | 80 +++++++++ .../integration/spend/test_cache_and_quota.py | 166 ++++++++++++++++++ 11 files changed, 626 insertions(+), 2 deletions(-) create mode 100644 tests/integration/_support/process.py create mode 100644 tests/integration/database/test_partition_transactions.py create mode 100644 tests/integration/database/test_reader_writer_regeneration.py create mode 100644 tests/integration/database/test_transaction_atomicity.py create mode 100644 tests/integration/pricing/test_price_precedence.py create mode 100644 tests/integration/spend/test_cache_and_quota.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 7c241af5853..3fcc2748115 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -2953,7 +2953,7 @@ workflows: name: integration-<< matrix.suite >> matrix: parameters: - suite: [management, accounting, providers] + suite: [management, accounting, database, providers] filters: branches: only: diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index b3bddfd898c..5c1ebe5d4f8 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -56,3 +56,7 @@ - {id: other.auth.jwt.wrong_audience_denied, module: other, tier: P0, area: auth, assertions: [wrong_audience_denied], source: "auth/handle_jwt.py", rationale: "A signed token from the trusted issuer intended for another app is rejected"} - {id: other.provider_wire.internal_parameters_filtered, module: other, tier: P0, area: provider_wire, assertions: [internal_parameters_filtered], source: "main.py", rationale: "A real provider request preserves content and public parameters without internal limiter fields"} - {id: other.provider_wire.validator_rejects_corruption, module: other, tier: P0, area: provider_wire, assertions: [validator_rejects_corruption], source: "tests/integration/_support/upstream.py", rationale: "The controlled transport rejects missing messages and internal fields while accepting supported metadata"} +- {id: other.database.partitions.lock_wait_outlives_transaction_default, module: other, tier: P0, area: database, assertions: [lock_wait_outlives_transaction_default], source: "spend_logs_partition_manager.py", rationale: "Actual partition DDL succeeds after a witnessed permitted lock wait beyond five seconds"} +- {id: other.database.partitions.repeat_preserves_rows, module: other, tier: P0, area: database, assertions: [repeat_preserves_rows], source: "spend_logs_partition_manager.py", rationale: "Repeated real partition maintenance preserves existing rows and one partition"} +- {id: other.database.regeneration.writer_updates_dependent_grants, module: other, tier: P0, area: database, assertions: [writer_updates_dependent_grants], source: "access_group_key_sync.py", rationale: "Regeneration preserves dependent grants when the configured reader rejects writes"} +- {id: other.database.access_group.failed_second_write_rolls_back_first, module: other, tier: P0, area: database, assertions: [failed_second_write_rolls_back_first], source: "access_group_endpoints.py", rationale: "A real constraint failure leaves neither group nor partial key grants"} diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 66198e6649c..0bc382a4744 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -66,3 +66,9 @@ - {id: quota_management.spend_tracking.custom_price.matches_input_rates, module: quota_management, tier: P0, behavior: spend_tracking, variant: custom_price, assertions: [matches_input_rates], exercised_on: [chat_completions], source: "router.py", rationale: "Configured deployment prices reach the response cost"} - {id: quota_management.spend_tracking.default_prices.survive_nullable_sibling_reload, module: quota_management, tier: P0, behavior: spend_tracking, variant: default_prices, assertions: [survive_nullable_sibling_reload], exercised_on: [chat_completions], source: "router.py", rationale: "Omitted and null prices retain defaults across sibling loading and reload, with persisted request charges"} - {id: quota_management.spend_tracking.default_prices.loaded_router_preserves_cached_defaults, module: quota_management, tier: P0, behavior: spend_tracking, variant: default_prices, assertions: [loaded_router_preserves_cached_defaults], exercised_on: [chat_completions], source: "router.py", rationale: "YAML-loaded omitted and null model-info prices preserve cached defaults across real SDK requests and reload order"} +- {id: quota_management.spend_tracking.price_precedence.zero_and_default_rates, module: quota_management, tier: P0, behavior: spend_tracking, variant: controlled_contract, exercised_on: [chat_completions], assertions: [zero_and_default_rates], source: "tests/integration", rationale: "Controlled requests with independent arithmetic, database state, and transport observations"} +- {id: quota_management.spend_tracking.alias_prices.remain_independent_on_reload, module: quota_management, tier: P0, behavior: spend_tracking, variant: controlled_contract, exercised_on: [chat_completions], assertions: [remain_independent_on_reload], source: "tests/integration", rationale: "Controlled requests with independent arithmetic, database state, and transport observations"} +- {id: quota_management.response_cache.generated_sequences_preserve_content_and_accounting, module: quota_management, tier: P0, behavior: spend_tracking, variant: controlled_contract, exercised_on: [chat_completions], assertions: [generated_sequences_preserve_content_and_accounting], source: "tests/integration", rationale: "Controlled requests with independent arithmetic, database state, and transport observations"} +- {id: quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores, module: quota_management, tier: P0, behavior: spend_tracking, variant: controlled_contract, exercised_on: [chat_completions], assertions: [boundary_blocks_before_provider_and_reset_restores], source: "tests/integration", rationale: "Controlled requests with independent arithmetic, database state, and transport observations"} +- {id: quota_management.response_cache.system_messages_partition_cache_identity, module: quota_management, tier: P0, behavior: spend_tracking, variant: controlled_contract, exercised_on: [chat_completions], assertions: [system_messages_partition_cache_identity], source: "tests/integration", rationale: "Controlled requests with independent arithmetic, database state, and transport observations"} +- {id: quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge, module: quota_management, tier: P0, behavior: spend_tracking, variant: repeated_cache_hit, exercised_on: [chat_completions], assertions: [single_charge], source: "spend_tracking_utils.py", rationale: "Repeated cached responses retain their identity while distinct spend rows bill only the original request"} diff --git a/tests/integration/README.md b/tests/integration/README.md index 3f918e51dcd..b0beb40e0a5 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -2,7 +2,7 @@ These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls -Use `tests/integration/run.py management`, `accounting` or `providers` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate +Use `tests/integration/run.py management`, `accounting`, `database` or `providers` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate Management also requires `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. CircleCI starts two directly addressed proxy processes sharing only that job's stores. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload @@ -17,3 +17,7 @@ Add contract definitions to the existing `tests/e2e/coverage_registry` and map c Provider sentinels currently use the controlled server, not live recordings. The provider shard also runs the existing strict replay controls for changed requests, exhausted interactions, leftover interactions and no provider connection. Future recorded scenarios must use that replay-only implementation; missing recordings cannot fall back to a real provider. The observation endpoint is destructive and the current selection runs serially against one owned upstream Fixtures must contain synthetic data only. Keep private incident records and source documents out of code, fixtures, logs and PR descriptions + +Database cases own their temporary schemas, roles, constraints and proxy processes. They prove reader-versus-writer execution with PostgreSQL lock observations, exercise real transaction wait limits and verify rollback after a reached database failure + +Accounting cases compare persisted input and output cost components against literal rates, including zero and default prices. Cache state models assert actual upstream calls, response identity and every persisted charge. Generated accounting tests have a 180-second test limit to accommodate the asynchronous spend writer; CircleCI keeps the whole shard capped at 11 minutes diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py new file mode 100644 index 00000000000..47c73825335 --- /dev/null +++ b/tests/integration/_support/process.py @@ -0,0 +1,92 @@ +import os +import socket +import signal +import subprocess +import sys +import time +import uuid +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from pathlib import Path +from typing import Final + +import httpx +import psutil + +from integration._support.client import Gateway + + +def in_group(process: psutil.Process, group: int) -> bool: + try: + return os.getpgid(process.pid) == group + except ProcessLookupError: + return False + + +def group_members(group: int) -> tuple[psutil.Process, ...]: + return tuple(process for process in psutil.process_iter() if in_group(process, group)) + + +def signal_group(group: int, action: int) -> None: + try: + os.killpg(group, action) + except ProcessLookupError: + pass + + +@contextmanager +def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str]) -> Iterator[Gateway]: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + port: Final = reserve.getsockname()[1] + root: Final = Path(__file__).resolve().parents[3] + environment: Final = { + **os.environ, + "LITELLM_MASTER_KEY": gateway.key, + "LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"), + "STORE_MODEL_IN_DB": "True", + **overrides, + } + output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) + output.mkdir(parents=True, exist_ok=True) + with (output / f"owned-proxy-{uuid.uuid4().hex}.log").open("w") as log: + process: Final = subprocess.Popen( + [sys.executable, "-m", "integration._support.proxy", "--config", "tests/integration/proxy_config.yaml", + "--host", "127.0.0.1", "--port", str(port), "--num_workers", "1", "--telemetry", "False", + "--use_prisma_db_push", "--enforce_prisma_migration_check"], + cwd=root, env=environment, stdout=log, stderr=subprocess.STDOUT, start_new_session=True, + ) + try: + with httpx.Client(base_url=f"http://127.0.0.1:{port}", timeout=15, trust_env=False) as client: + deadline: Final = time.monotonic() + 70 + while True: + assert process.poll() is None, "Owned proxy exited before readiness" + try: + if client.get("/health/readiness", timeout=2).status_code == 200: + break + except httpx.TransportError: + pass + assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded" + time.sleep(0.1) + yield Gateway(client, gateway.key, gateway.upstream_url) + finally: + forced = False + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + forced = True + residual: Final = group_members(process.pid) + if residual: + signal_group(process.pid, signal.SIGTERM) + psutil.wait_procs(residual, timeout=5) + remaining: Final = group_members(process.pid) + if remaining: + forced = True + signal_group(process.pid, signal.SIGKILL) + psutil.wait_procs(remaining, timeout=3) + process.wait(timeout=3) + survivors: Final = group_members(process.pid) + assert not survivors, "Owned proxy child survived cleanup" + assert not forced, "Owned proxy required forced cleanup" diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 82cc64dd5c6..8c09048c243 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -75,6 +75,34 @@ ], "tests/integration/authorization/test_warmed_policy.py::test_expiry_and_explicit_clear_reach_both_warmed_workers": [ "mgmt.key.update.expiry_changes_reach_warmed_workers" + ], + "tests/integration/database/test_partition_transactions.py::test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent": [ + "other.database.partitions.lock_wait_outlives_transaction_default", + "other.database.partitions.repeat_preserves_rows" + ], + "tests/integration/database/test_reader_writer_regeneration.py::test_key_regeneration_uses_writer_with_a_real_readonly_reader": [ + "other.database.regeneration.writer_updates_dependent_grants" + ], + "tests/integration/pricing/test_price_precedence.py::test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic": [ + "quota_management.spend_tracking.price_precedence.zero_and_default_rates" + ], + "tests/integration/pricing/test_price_precedence.py::test_same_upstream_aliases_keep_distinct_prices_after_reload": [ + "quota_management.spend_tracking.alias_prices.remain_independent_on_reload" + ], + "tests/integration/spend/test_cache_and_quota.py::test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost": [ + "quota_management.response_cache.generated_sequences_preserve_content_and_accounting" + ], + "tests/integration/spend/test_cache_and_quota.py::test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores": [ + "quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores" + ], + "tests/integration/spend/test_cache_and_quota.py::test_different_system_messages_do_not_share_a_cached_response": [ + "quota_management.response_cache.system_messages_partition_cache_identity" + ], + "tests/integration/database/test_transaction_atomicity.py::test_access_group_second_key_constraint_failure_rolls_back_all_writes": [ + "other.database.access_group.failed_second_write_rolls_back_first" + ], + "tests/integration/spend/test_cache_and_quota.py::test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows": [ + "quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge" ] } } diff --git a/tests/integration/database/test_partition_transactions.py b/tests/integration/database/test_partition_transactions.py new file mode 100644 index 00000000000..51238458f71 --- /dev/null +++ b/tests/integration/database/test_partition_transactions.py @@ -0,0 +1,86 @@ +import asyncio +import os +import time +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Final +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +import psycopg +import pytest +from psycopg import sql +from prisma import Prisma + +from integration._support.database import read_rows +from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import SpendLogsPartitionManager + + +@dataclass(frozen=True) +class PartitionConnection: + db: Prisma + + +@pytest.mark.covers("other.database.partitions.lock_wait_outlives_transaction_default", "other.database.partitions.repeat_preserves_rows") +async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> None: + schema: Final = f"integration_{uuid.uuid4().hex}" + url: Final = os.environ["DATABASE_URL"] + parsed: Final = urlsplit(url) + scoped_url: Final = urlunsplit(parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema}))) + parent: Final = sql.Identifier(schema, "LiteLLM_SpendLogs") + with psycopg.connect(url, autocommit=True) as setup: + setup.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema))) + try: + setup.execute(sql.SQL('CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL) PARTITION BY RANGE ("startTime")').format(parent)) + database: Final = Prisma(datasource={"url": scoped_url}) + await database.connect() + try: + manager: Final = SpendLogsPartitionManager(interval="day", precreate_ahead=0) + with psycopg.connect(url) as blocker: + blocker.execute(sql.SQL("LOCK TABLE {} IN ACCESS SHARE MODE").format(parent)) + blocker_pid: Final = blocker.info.backend_pid + operation: Final = asyncio.create_task(manager.ensure_partitions(PartitionConnection(database), lambda: 7000)) + wait_deadline: Final = time.monotonic() + 3 + try: + while True: + witnesses: Final = read_rows( + "SELECT a.pid, extract(epoch FROM clock_timestamp()-a.query_start)::double precision AS age " + "FROM pg_stat_activity a WHERE %s = ANY(pg_blocking_pids(a.pid)) " + "AND a.wait_event_type = 'Lock' AND a.query LIKE 'CREATE TABLE IF NOT EXISTS%%'", + (blocker_pid,), + ) + if witnesses: + break + assert time.monotonic() < wait_deadline, "Partition DDL never reached the held lock" + await asyncio.sleep(0.02) + assert len(witnesses) == 1 + held_at: Final = time.monotonic() + age: Final = float(witnesses[0]["age"]) + assert age < 1, "DDL lock witness arrived too late for the qualification window" + await asyncio.sleep(5.6 - age) + held_seconds: Final = age + time.monotonic() - held_at + assert 5.5 <= held_seconds < 6.5, f"Lock qualification timing outside window: {held_seconds}" + assert not operation.done(), "DDL completed while its required lock was held" + except BaseException: + operation.cancel() + await asyncio.gather(operation, return_exceptions=True) + raise + finally: + blocker.rollback() + ensured: Final = await asyncio.wait_for(operation, timeout=5) + assert len(ensured) == 1, "Partition DDL failed after the permitted lock wait" + catalog: Final = read_rows( + "SELECT child.relname FROM pg_inherits i JOIN pg_class child ON child.oid=i.inhrelid " + "JOIN pg_class parent ON parent.oid=i.inhparent JOIN pg_namespace n ON n.oid=parent.relnamespace " + "WHERE n.nspname=%s AND parent.relname='LiteLLM_SpendLogs'", (schema,), + ) + assert catalog == [{"relname": ensured[0]}] + now: Final = datetime.now(timezone.utc).replace(tzinfo=None) + setup.execute(sql.SQL('INSERT INTO {} VALUES (%s, %s)').format(parent), ("retained", now)) + assert await manager.ensure_partitions(PartitionConnection(database), lambda: 7000) == ensured + assert setup.execute(sql.SQL("SELECT request_id FROM {}").format(parent)).fetchall() == [("retained",)] + finally: + await database.disconnect() + finally: + setup.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema))) + assert read_rows("SELECT nspname FROM pg_namespace WHERE nspname=%s", (schema,)) == [] diff --git a/tests/integration/database/test_reader_writer_regeneration.py b/tests/integration/database/test_reader_writer_regeneration.py new file mode 100644 index 00000000000..382dd14fbb0 --- /dev/null +++ b/tests/integration/database/test_reader_writer_regeneration.py @@ -0,0 +1,91 @@ +import os +import uuid +from hashlib import sha256 +from pathlib import Path +from concurrent.futures import ThreadPoolExecutor +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +import psycopg +import pytest +from psycopg import sql + +from integration._support.client import Gateway, eventually, string_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy + + +def delete_if_present(candidate: Gateway, key: str) -> None: + digest: Final = sha256(key.encode()).hexdigest() + if read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)): + candidate.post("/key/delete", {"keys": [key]}) + assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [] + + +@pytest.mark.covers("other.database.regeneration.writer_updates_dependent_grants") +def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gateway, tmp_path: Path) -> None: + role: Final = f"integration_reader_{uuid.uuid4().hex}" + url: Final = os.environ["DATABASE_URL"] + parsed: Final = urlsplit(url) + reader_url: Final = urlunsplit(parsed._replace(netloc=f"{role}:integration-reader-password@{parsed.hostname}:{parsed.port}")) + with psycopg.connect(url, autocommit=True) as admin: + admin.execute(sql.SQL("CREATE ROLE {} LOGIN PASSWORD 'integration-reader-password' NOSUPERUSER NOINHERIT").format(sql.Identifier(role))) + try: + admin.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(role))) + admin.execute(sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA public TO {}").format(sql.Identifier(role))) + admin.execute(sql.SQL("ALTER ROLE {} SET default_transaction_read_only = on").format(sql.Identifier(role))) + with psycopg.connect(reader_url, autocommit=True) as reader: + assert reader.execute("SHOW transaction_read_only").fetchone() == ("on",) + with pytest.raises(psycopg.errors.ReadOnlySqlTransaction): + reader.execute('UPDATE "LiteLLM_VerificationToken" SET blocked = true WHERE false') + with owned_proxy(gateway, tmp_path, {"DATABASE_URL_READ_REPLICA": reader_url}) as candidate: + assert read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,)), "Candidate reader was never connected" + with gateway.scenario() as scenario: + model: Final = scenario.model() + outside: Final = scenario.model() + old: Final = string_value(candidate.post("/key/generate", {"models": [outside]})["key"]) + new: Final = f"sk-integration-{uuid.uuid4().hex}" + scenario.cleanups.callback(delete_if_present, gateway, old) + scenario.cleanups.callback(delete_if_present, gateway, new) + old_hash: Final = sha256(old.encode()).hexdigest() + before: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "no grant yet"}]}, key=old) + assert before.status_code == 403 and before.json()["error"]["type"] == "key_model_access_denied", before.text + response: Final = candidate.request("POST", "/v1/access_group", { + "access_group_name": f"integration-{uuid.uuid4().hex}", + "access_model_names": [model], "assigned_key_ids": [old_hash], + }) + assert response.status_code == 201, response.text + group: Final = string_value(response.json()["access_group_id"]) + try: + with psycopg.connect(url) as blocker, ThreadPoolExecutor(max_workers=1) as executor: + blocker.execute('LOCK TABLE "LiteLLM_AccessGroupTable" IN ACCESS EXCLUSIVE MODE') + pending: Final = executor.submit(candidate.request, "GET", f"/v1/access_group/{group}") + try: + reached: Final = eventually(lambda: read_rows( + "SELECT usename FROM pg_stat_activity WHERE %s=ANY(pg_blocking_pids(pid)) " + "AND usename=%s AND query LIKE 'SELECT%%'", (blocker.info.backend_pid, role), + ), bool, seconds=3) + assert reached == [{"usename": role}] + finally: + blocker.rollback() + selected: Final = pending.result(timeout=5) + assert selected.status_code == 200 and selected.json()["access_group_id"] == group, selected.text + assert candidate.chat(model, key=old)["usage"]["total_tokens"] == 40 + regenerated: Final = candidate.post("/key/regenerate", {"key": old, "new_key": new, "grace_period": "0s"}) + assert regenerated["key"] == new + new_hash: Final = sha256(new.encode()).hexdigest() + assert new != old + assert read_rows('SELECT assigned_key_ids FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)) == [{"assigned_key_ids": [new_hash]}] + assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s)', ([old_hash, new_hash],)) == [{"token": new_hash, "access_group_ids": [group]}] + assert candidate.chat(model, key=new)["usage"]["total_tokens"] == 40 + assert candidate.chat(outside, key=new)["usage"]["total_tokens"] == 40 + denied: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rotated key"}]}, key=old) + assert denied.status_code == 401 and denied.json()["error"]["type"] == "token_not_found_in_db", denied.text + finally: + deleted: Final = gateway.request("DELETE", f"/v1/access_group/{group}") + assert deleted.status_code == 204, deleted.text + assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)) == [] + finally: + admin.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role))) + admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role))) + assert read_rows("SELECT rolname FROM pg_roles WHERE rolname=%s", (role,)) == [] diff --git a/tests/integration/database/test_transaction_atomicity.py b/tests/integration/database/test_transaction_atomicity.py new file mode 100644 index 00000000000..37b5d2289b7 --- /dev/null +++ b/tests/integration/database/test_transaction_atomicity.py @@ -0,0 +1,67 @@ +import os +import uuid +from contextlib import ExitStack +from hashlib import sha256 +from typing import Final + +import psycopg +import pytest +from psycopg import sql + +from integration._support.client import Gateway +from integration._support.database import read_rows + + +@pytest.mark.covers("other.database.access_group.failed_second_write_rolls_back_first") +def test_access_group_second_key_constraint_failure_rolls_back_all_writes(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model() + outside: Final = scenario.model() + keys: Final = (scenario.key(models=[outside]), scenario.key(models=[outside])) + tokens: Final = [sha256(key.encode()).hexdigest() for key in keys] + name: Final = f"integration-{uuid.uuid4().hex}" + constraint: Final = f"integration_reject_{uuid.uuid4().hex}" + witness: Final = constraint + "_seq" + check_function: Final = constraint + "_check" + body: Final = {"access_group_name": name, "access_model_names": [model], "assigned_key_ids": tokens} + def remove_partial_group() -> None: + for row in read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)): + response: Final = gateway.request("DELETE", f"/v1/access_group/{row['access_group_id']}") + assert response.status_code == 204, response.text + assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)) == [] + + scenario.cleanups.callback(remove_partial_group) + before: Final = read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection, ExitStack() as cleanup: + connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(witness))) + cleanup.callback(connection.execute, sql.SQL("DROP SEQUENCE {}").format(sql.Identifier(witness))) + connection.execute(sql.SQL("CREATE FUNCTION {}(text[]) RETURNS boolean LANGUAGE plpgsql AS $$ BEGIN IF cardinality($1)>0 THEN PERFORM nextval({}); RETURN false; END IF; RETURN true; END $$").format(sql.Identifier(check_function), sql.Literal(witness))) + cleanup.callback(connection.execute, sql.SQL("DROP FUNCTION {}(text[])").format(sql.Identifier(check_function))) + connection.execute(sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT {} CHECK (token <> {} OR {}(access_group_ids))').format(sql.Identifier(constraint), sql.Literal(tokens[1]), sql.Identifier(check_function))) + cleanup.callback(connection.execute, sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" DROP CONSTRAINT {}').format(sql.Identifier(constraint))) + try: + assert connection.execute(sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness))).fetchone() == (False,) + failed: Final = gateway.request("POST", "/v1/access_group", body) + assert failed.status_code == 500, failed.text + # Sequence advancement survives rollback and proves the rejecting + # constraint actually evaluated the second key's nonempty grant. + assert connection.execute(sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness))).fetchone() == (True,) + assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)) == [] + assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) == before + for key in keys: + denied: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rolled back grant"}]}, key=key) + assert denied.status_code == 403 and denied.json()["error"]["type"] == "key_model_access_denied", denied.text + finally: + cleanup.close() + created: Final = gateway.request("POST", "/v1/access_group", body) + assert created.status_code == 201, created.text + identity: Final = created.json()["access_group_id"] + try: + for key in keys: + assert gateway.chat(model, key=key)["usage"]["total_tokens"] == 40 + finally: + deleted: Final = gateway.request("DELETE", f"/v1/access_group/{identity}") + assert deleted.status_code == 204, deleted.text + assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (identity,)) == [] + assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) == before + assert read_rows('SELECT conname FROM pg_constraint WHERE conname=%s', (constraint,)) == [] diff --git a/tests/integration/pricing/test_price_precedence.py b/tests/integration/pricing/test_price_precedence.py new file mode 100644 index 00000000000..aa8e8628cd8 --- /dev/null +++ b/tests/integration/pricing/test_price_precedence.py @@ -0,0 +1,80 @@ +import json +import uuid +from typing import Final + +import pytest +from hypothesis import Phase, example, given, settings, strategies as st + +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows + + +@pytest.mark.covers("quota_management.spend_tracking.price_precedence.zero_and_default_rates") +@pytest.mark.timeout(180) +def test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic(gateway: Gateway) -> None: + @settings(max_examples=20, deadline=None, database=None, phases=(Phase.explicit, Phase.generate, Phase.shrink)) + @example(rates=(0, 0)) + @example(rates=(1, 2)) + @example(rates=("null", "null")) + @given(rates=st.one_of(st.sampled_from((("omitted", "omitted"), ("null", "null"))), st.tuples(st.integers(0, 25), st.integers(0, 25)))) + def check(rates: tuple[str | int, str | int]) -> None: + if rates[0] in ("omitted", "null"): + parameters: Final = {} if rates[0] == "omitted" else {"input_cost_per_token": None, "output_cost_per_token": None} + input_rate, output_rate = 0.00000015, 0.0000006 + else: + assert isinstance(rates[0], int) and isinstance(rates[1], int) + input_rate, output_rate = rates[0] / 1_000_000, rates[1] / 1_000_000 + parameters = {"input_cost_per_token": input_rate, "output_cost_per_token": output_rate} + with gateway.scenario() as scenario: + model: Final = scenario.model(**parameters) + response: Final = gateway.request("POST", "/v1/chat/completions", { + "model": model, "messages": [{"role": "user", "content": f"independent price {uuid.uuid4().hex}"}], + }) + assert response.status_code == 200, response.text + assert response.json()["usage"] == {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40} + expected: Final = 20 * input_rate + 20 * output_rate + if expected: + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6) + else: + assert response.headers.get("x-litellm-response-cost") in (None, "0", "0.0") + rows: Final = eventually(lambda: read_rows( + 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (response.json()["id"],), + ), lambda values: len(values) == 1, seconds=70) + assert rows[0]["prompt_tokens"] == 20 and rows[0]["completion_tokens"] == 20 + assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6) + metadata: Final = rows[0]["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + breakdown: Final = object_value(parsed["cost_breakdown"]) + assert float(breakdown["input_cost"]) == pytest.approx(20 * input_rate, rel=1e-6) + assert float(breakdown["output_cost"]) == pytest.approx(20 * output_rate, rel=1e-6) + + check() + + +@pytest.mark.covers("quota_management.spend_tracking.alias_prices.remain_independent_on_reload") +def test_same_upstream_aliases_keep_distinct_prices_after_reload(gateway: Gateway) -> None: + for order in (("free", "paid"), ("paid", "free")): + with gateway.scenario() as scenario: + rates: Final = {"free": {"input_cost_per_token": 0, "output_cost_per_token": 0}, "paid": {"input_cost_per_token": 0.001, "output_cost_per_token": 0.003}} + aliases: Final = {kind: scenario.model(**rates[kind]) for kind in order} + for generation in range(2): + for kind in order if generation == 0 else reversed(order): + model: Final = aliases[kind] + cost: Final = 0.08 if kind == "paid" else 0.0 + response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"alias price {model} {generation}"}]}) + assert response.status_code == 200, response.text + assert response.json()["usage"]["total_tokens"] == 40 + if cost: + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(cost) + rows: Final = eventually(lambda response=response: read_rows('SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response.json()["id"],)), lambda values: len(values) == 1, seconds=70) + assert float(rows[0]["spend"]) == pytest.approx(cost) + metadata: Final = rows[0]["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + breakdown: Final = object_value(parsed["cost_breakdown"]) + assert float(breakdown["input_cost"]) == pytest.approx(20 * rates[kind]["input_cost_per_token"]) + assert float(breakdown["output_cost"]) == pytest.approx(20 * rates[kind]["output_cost_per_token"]) + entries: Final = gateway.get("/model/info")["data"] + target: Final = next(entry for entry in entries if entry["model_name"] == aliases["paid"]) + changed: Final = gateway.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "price reload"}}) + assert changed.status_code == 200, changed.text diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py new file mode 100644 index 00000000000..2e66a538161 --- /dev/null +++ b/tests/integration/spend/test_cache_and_quota.py @@ -0,0 +1,166 @@ +import uuid +from contextlib import ExitStack +from hashlib import sha256 +from typing import Final + +import httpx +import pytest +from hypothesis import strategies as st +from hypothesis.stateful import RuleBasedStateMachine, rule, run_state_machine_as_test + +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests + + +@pytest.mark.covers("quota_management.response_cache.generated_sequences_preserve_content_and_accounting") +@pytest.mark.timeout(180) +def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gateway: Gateway) -> None: + class CacheRequests(RuleBasedStateMachine): + def __init__(self) -> None: + super().__init__() + self.resources = ExitStack() + try: + self.scenario = self.resources.enter_context(gateway.scenario()) + self.upstream = self.resources.enter_context(httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False)) + self.model = self.scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + self.key = self.scenario.key(models=[self.model]) + self.prefix = uuid.uuid4().hex + self.seen: frozenset[int] = frozenset() + self.requests = 0 + self.paid = 0 + self.failed = False + self.identities: dict[int, str] = {} + except BaseException: + with budget.cleanup(): + self.resources.close() + raise + + @rule(marker=st.integers(min_value=0, max_value=2)) + def request(self, marker: int) -> None: + try: + self.perform_request(marker) + except BaseException: + self.failed = True + raise + + def perform_request(self, marker: int) -> None: + self.upstream.get("/__observations").raise_for_status() + response: Final = gateway.request("POST", "/v1/chat/completions", { + "model": self.model, "messages": [{"role": "user", "content": f"{self.prefix}-{marker}"}], + }, key=self.key) + assert response.status_code == 200, response.text + self.requests += 1 + body: Final = response.json() + assert body["choices"][0]["message"]["content"] == "Hello! This is a mock response from the fake OpenAI endpoint." + assert body["usage"]["total_tokens"] == 40 + observed: Final = self.upstream.get("/__observations").json()["requests"] + expected_calls: Final = 0 if marker in self.seen else 1 + assert len(observed) == expected_calls, observed + # The response can retain its original cost header on a cache hit. + # Per-request billed cost is checked against fresh spend rows below. + if marker not in self.seen: + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(0.06) + if marker in self.identities: + assert body["id"] == self.identities[marker] + else: + assert body["id"] not in self.identities.values() + self.identities = {**self.identities, marker: body["id"]} + self.paid += expected_calls + self.seen = self.seen.union((marker,)) + + def teardown(self) -> None: + try: + if self.requests and not self.failed: + rows: Final = eventually(lambda: read_rows( + 'SELECT request_id, spend, cache_hit, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(self.key.encode()).hexdigest(),), + ), lambda values: len(values) == self.requests, seconds=70) + assert len({row["request_id"] for row in rows}) == self.requests + assert sum(float(row["spend"]) for row in rows) == pytest.approx(self.paid * 0.06) + assert sum(row["cache_hit"] == "True" for row in rows) == self.requests - self.paid + for row in rows: + assert row["prompt_tokens"] == 20 and row["completion_tokens"] == 20 + if row["cache_hit"] == "True": + assert float(row["spend"]) == 0 and "_cache_hit" in row["request_id"] + assert any(row["request_id"].startswith(identity + "_cache_hit") for identity in self.identities.values()) + else: + assert row["request_id"] in self.identities.values() + assert float(row["spend"]) == pytest.approx(0.06) + finally: + with budget.cleanup(): + self.resources.close() + + with bounded_http_requests((gateway,), limit=2000) as budget: + run_state_machine_as_test(CacheRequests, settings=LIFECYCLE_SETTINGS) + + +@pytest.mark.covers("quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge") +def test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows(gateway: Gateway) -> None: + with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model]) + prompt: Final = f"repeated cache {uuid.uuid4().hex}" + upstream.get("/__observations").raise_for_status() + results: Final = tuple(gateway.chat(model, key=key, text=prompt) for _ in range(3)) + assert len(upstream.get("/__observations").json()["requests"]) == 1 + assert len({result["id"] for result in results}) == 1 + for result in results: + assert result["choices"][0]["message"]["content"] == "Hello! This is a mock response from the fake OpenAI endpoint." + assert result["usage"]["total_tokens"] == 40 + rows: Final = eventually(lambda: read_rows('SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)), lambda values: len(values) == 3, seconds=70) + assert len({row["request_id"] for row in rows}) == 3 + assert sorted(float(row["spend"]) for row in rows) == [0, 0, 0.06] + for row in rows: + if row["cache_hit"] == "True": + assert float(row["spend"]) == 0 + assert row["request_id"].startswith(results[0]["id"] + "_cache_hit") + else: + assert row["request_id"] == results[0]["id"] and float(row["spend"]) == pytest.approx(0.06) + + +@pytest.mark.covers("quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores") +def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gateway: Gateway) -> None: + with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model], max_budget=0.06) + control: Final = scenario.key(models=[model]) + first: Final = gateway.chat(model, key=key, text=f"budget {uuid.uuid4().hex}") + assert first["usage"]["total_tokens"] == 40 + digest: Final = sha256(key.encode()).hexdigest() + spent: Final = eventually(lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, seconds=70) + assert float(spent[0]["spend"]) == pytest.approx(0.06) + upstream.get("/__observations").raise_for_status() + denied: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"over budget {uuid.uuid4().hex}"}]}, key=key) + assert denied.status_code == 429 and denied.json()["error"]["type"] == "budget_exceeded", denied.text + assert upstream.get("/__observations").json()["requests"] == [] + assert gateway.chat(model, key=control, text=f"control {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + gateway.post("/key/update", {"key": key, "spend": 0}) + assert read_rows('SELECT spend, max_budget FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [{"spend": 0.0, "max_budget": 0.06}] + assert gateway.chat(model, key=key, text=f"reset {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + eventually(lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, seconds=70) + upstream.get("/__observations").raise_for_status() + denied_again: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"boundary again {uuid.uuid4().hex}"}]}, key=key) + assert denied_again.status_code == 429 and denied_again.json()["error"]["type"] == "budget_exceeded", denied_again.text + assert upstream.get("/__observations").json()["requests"] == [] + + +@pytest.mark.covers("quota_management.response_cache.system_messages_partition_cache_identity") +def test_different_system_messages_do_not_share_a_cached_response(gateway: Gateway) -> None: + with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + model: Final = scenario.model() + prompt: Final = uuid.uuid4().hex + identities: dict[str, str] = {} + for system, expected_calls in (("first policy", 1), ("second policy", 1), ("first policy", 0)): + upstream.get("/__observations").raise_for_status() + response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "system", "content": system}, {"role": "user", "content": prompt}]}) + assert response.status_code == 200 and response.json()["usage"]["total_tokens"] == 40, response.text + calls: Final = upstream.get("/__observations").json()["requests"] + assert len(calls) == expected_calls + if system in identities: + assert response.json()["id"] == identities[system] + else: + assert response.json()["id"] not in identities.values() + identities = {**identities, system: response.json()["id"]} + if calls: + assert calls[0]["body"]["messages"] == [{"role": "system", "content": system}, {"role": "user", "content": prompt}] From 57825775b830fdf479d34b1b58e113c4f107083c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 05:08:14 -0700 Subject: [PATCH 022/187] test: tighten reload coverage and avoid an artificial lock timing cutoff --- .../integration/database/test_partition_transactions.py | 5 ++--- tests/integration/pricing/test_price_precedence.py | 9 +++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/integration/database/test_partition_transactions.py b/tests/integration/database/test_partition_transactions.py index 51238458f71..2759d8dcfb7 100644 --- a/tests/integration/database/test_partition_transactions.py +++ b/tests/integration/database/test_partition_transactions.py @@ -56,10 +56,9 @@ async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> assert len(witnesses) == 1 held_at: Final = time.monotonic() age: Final = float(witnesses[0]["age"]) - assert age < 1, "DDL lock witness arrived too late for the qualification window" - await asyncio.sleep(5.6 - age) + await asyncio.sleep(max(0, 5.6 - age)) held_seconds: Final = age + time.monotonic() - held_at - assert 5.5 <= held_seconds < 6.5, f"Lock qualification timing outside window: {held_seconds}" + assert held_seconds >= 5.5, f"Lock released before the transaction boundary: {held_seconds}" assert not operation.done(), "DDL completed while its required lock was held" except BaseException: operation.cancel() diff --git a/tests/integration/pricing/test_price_precedence.py b/tests/integration/pricing/test_price_precedence.py index aa8e8628cd8..1f724270a81 100644 --- a/tests/integration/pricing/test_price_precedence.py +++ b/tests/integration/pricing/test_price_precedence.py @@ -74,7 +74,8 @@ def test_same_upstream_aliases_keep_distinct_prices_after_reload(gateway: Gatewa breakdown: Final = object_value(parsed["cost_breakdown"]) assert float(breakdown["input_cost"]) == pytest.approx(20 * rates[kind]["input_cost_per_token"]) assert float(breakdown["output_cost"]) == pytest.approx(20 * rates[kind]["output_cost_per_token"]) - entries: Final = gateway.get("/model/info")["data"] - target: Final = next(entry for entry in entries if entry["model_name"] == aliases["paid"]) - changed: Final = gateway.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "price reload"}}) - assert changed.status_code == 200, changed.text + if generation == 0: + entries: Final = gateway.get("/model/info")["data"] + target: Final = next(entry for entry in entries if entry["model_name"] == aliases["paid"]) + changed: Final = gateway.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "price reload"}}) + assert changed.status_code == 200, changed.text From 9fbc2c5b71251d879d02cb8269fd0191fe7ec448 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 16:23:06 +0000 Subject: [PATCH 023/187] fix(proxy): show all model groups to proxy admins in /model_group/info Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 32 +++++---- .../proxy_server/test_routes_model_info.py | 70 ++++++++++++++++--- 2 files changed, 79 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..841ac47399e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15494,18 +15494,26 @@ async def model_group_info( from litellm.proxy.utils import get_available_models_for_user # Get available models for the user - all_models_str: Final = await get_available_models_for_user( - user_api_key_dict=user_api_key_dict, - llm_router=llm_router, - general_settings=general_settings, - user_model=user_model, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - team_id=None, - include_model_access_groups=False, - only_model_access_groups=False, - return_wildcard_routes=False, - user_api_key_cache=user_api_key_cache, + is_proxy_admin: Final = user_api_key_dict.user_role in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ) + all_models_str: Final = ( + llm_router.get_model_names() + if is_proxy_admin + else await get_available_models_for_user( + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + general_settings=general_settings, + user_model=user_model, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + team_id=None, + include_model_access_groups=False, + only_model_access_groups=False, + return_wildcard_routes=False, + user_api_key_cache=user_api_key_cache, + ) ) model_groups: list[ModelGroupInfoProxy] = _get_model_group_info( llm_router=llm_router, all_models_str=all_models_str, model_group=model_group diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 4c141bcf698..a076b3593b1 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -9,7 +9,7 @@ Pins (PR2): from __future__ import annotations -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest @@ -128,7 +128,6 @@ def test_v1_model_info_no_model_list_error(client, auth_as, null_router, path): assert "LLM Model List not loaded" in response.text - def test_get_proxy_model_info_surfaces_supports_parallel_function_calling(local_model_cost_map): """``GET /v1/model/info`` enriches each deployment through ``_get_proxy_model_info``; a registry entry declaring parallel function calling must land in ``model_info`` instead of null.""" @@ -161,9 +160,7 @@ def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monkeypatch router.get_model_list = MagicMock(return_value=[deployment]) monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) - expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info( - [deployment] - ) + expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info([deployment]) allowed_model_names = proxy_server._get_v1_model_info_allowed_model_names( user_api_key_dict=UserAPIKeyAuth( api_key="sk-test", @@ -308,6 +305,61 @@ def test_model_group_info_invalid_method(client, auth_as, null_router): assert len(response.content) > 0 +@pytest.fixture +def model_group_info_router(monkeypatch): + from litellm.types.proxy.management_endpoints.model_management_endpoints import ModelGroupInfoProxy + + model_names = ["gpt-4", "claude-3"] + router = MagicMock() + router.get_model_names.return_value = model_names + router.get_model_access_groups.return_value = {} + router.get_model_list.return_value = [] + + def model_group_info(*, llm_router, all_models_str, model_group): + return [ModelGroupInfoProxy(model_group=name, providers=[]) for name in all_models_str] + + async def append_agents_to_model_group(*, model_groups, user_api_key_dict): + return model_groups + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", [{"model_name": name} for name in model_names]) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", None) + monkeypatch.setattr(proxy_server, "user_api_key_cache", None) + monkeypatch.setattr(proxy_server, "_get_model_group_info", model_group_info) + + from litellm.proxy.agent_endpoints import model_list_helpers + + monkeypatch.setattr( + model_list_helpers, + "append_agents_to_model_group", + AsyncMock(side_effect=append_agents_to_model_group), + ) + return router + + +def test_model_group_info_proxy_admin_ignores_key_model_restriction(client, auth_as, model_group_info_router): + from litellm.proxy._types import LitellmUserRoles + + with auth_as(LitellmUserRoles.PROXY_ADMIN, models=["no-default-models"]): + response = client.get("/model_group/info") + + assert response.status_code == 200 + assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4", "claude-3"] + + +def test_model_group_info_internal_user_key_model_restriction_applies(client, auth_as, model_group_info_router): + from litellm.proxy._types import LitellmUserRoles + + with auth_as(LitellmUserRoles.INTERNAL_USER, models=["gpt-4"]): + response = client.get("/model_group/info") + + assert response.status_code == 200 + assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4"] + + # --------------------------------------------------------------------------- # GET /v2/model/info?exclude_auto_routers # --------------------------------------------------------------------------- @@ -399,14 +451,10 @@ def test_v2_model_info_exclude_auto_routers_shrinks_total_count(client, auth_as, assert len(payload["data"]) == payload["total_count"] -def test_v2_model_info_exclude_auto_routers_paginates_over_the_filtered_set( - client, auth_as, mixed_auto_router_router -): +def test_v2_model_info_exclude_auto_routers_paginates_over_the_filtered_set(client, auth_as, mixed_auto_router_router): """Page size applies to the filtered list, so no page silently comes back short.""" with auth_as(): - response = client.get( - "/v2/model/info", params={"exclude_auto_routers": "true", "page": 1, "size": 1} - ) + response = client.get("/v2/model/info", params={"exclude_auto_routers": "true", "page": 1, "size": 1}) payload = response.json() assert payload["total_count"] == 2 assert payload["total_pages"] == 2 From 834f187f4c9ea34be582b2e7315b93665124d9c0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 19:20:36 +0000 Subject: [PATCH 024/187] feat(models): add Together, Fireworks, Nebius, Cohere, Bedrock models and Inception cache pricing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 88 +++++++++++++++++++ model_prices_and_context_window.json | 88 +++++++++++++++++++ 2 files changed, 176 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 89d582db151..0954f9360fa 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -352,6 +352,17 @@ "supports_function_calling": true, "supports_pdf_input": true }, + "writer.palmyra-vision-7b": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-writer-palmyra-vision-7b.html", + "supports_vision": true + }, "amazon.nova-lite-v1:0": { "input_cost_per_token": 6e-08, "litellm_provider": "bedrock_converse", @@ -14650,6 +14661,21 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "command-a-plus-05-2026": { + "input_cost_per_token": 0.0, + "litellm_provider": "cohere_chat", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.cohere.com/docs/command-a-plus", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "command-light": { "input_cost_per_token": 3e-07, "litellm_provider": "cohere_chat", @@ -34159,6 +34185,7 @@ "supports_tool_choice": true }, "inception/mercury-2.5": { + "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "inception", "max_input_tokens": 260000, @@ -34168,6 +34195,7 @@ "output_cost_per_token": 7.5e-07, "source": "https://docs.inceptionlabs.ai/get-started/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true @@ -37333,6 +37361,15 @@ "supports_reasoning": true, "source": "https://tokenfactory.nebius.com/models/catalog/text2text/deepseek-ai%2FDeepSeek-V4-Pro" }, + "nebius/deepseek-ai/DeepSeek-V4-Pro-0813": { + "input_cost_per_token": 1.32e-06, + "litellm_provider": "nebius", + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/deepseek-ai%2FDeepSeek-V4-Pro-0813", + "supports_function_calling": true, + "supports_reasoning": true + }, "nebius/MiniMaxAI/MiniMax-M2.5": { "max_tokens": 196608, "max_input_tokens": 196608, @@ -37599,6 +37636,17 @@ "supports_reasoning": true, "source": "https://tokenfactory.nebius.com/models/catalog/text2text/zai-org%2FGLM-5.2" }, + "nebius/zai-org/GLM-5.3": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "nebius", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/zai-org%2FGLM-5.3", + "supports_function_calling": true, + "supports_reasoning": true + }, "nebius/zai-org/GLM-5.3-Flash": { "max_tokens": 1024000, "max_input_tokens": 1024000, @@ -44017,6 +44065,20 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "together_ai/deepseek-ai/DeepSeek-V4-Pro": { "deprecation_date": "2026-08-27", "cache_read_input_token_cost": 2e-07, @@ -61060,6 +61122,32 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/glm-5p3-fast": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/models/glm-5p3-flash": { "cache_read_input_token_cost": 3e-08, "input_cost_per_token": 1.5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 89d582db151..0954f9360fa 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -352,6 +352,17 @@ "supports_function_calling": true, "supports_pdf_input": true }, + "writer.palmyra-vision-7b": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-writer-palmyra-vision-7b.html", + "supports_vision": true + }, "amazon.nova-lite-v1:0": { "input_cost_per_token": 6e-08, "litellm_provider": "bedrock_converse", @@ -14650,6 +14661,21 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "command-a-plus-05-2026": { + "input_cost_per_token": 0.0, + "litellm_provider": "cohere_chat", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.cohere.com/docs/command-a-plus", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "command-light": { "input_cost_per_token": 3e-07, "litellm_provider": "cohere_chat", @@ -34159,6 +34185,7 @@ "supports_tool_choice": true }, "inception/mercury-2.5": { + "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "inception", "max_input_tokens": 260000, @@ -34168,6 +34195,7 @@ "output_cost_per_token": 7.5e-07, "source": "https://docs.inceptionlabs.ai/get-started/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true @@ -37333,6 +37361,15 @@ "supports_reasoning": true, "source": "https://tokenfactory.nebius.com/models/catalog/text2text/deepseek-ai%2FDeepSeek-V4-Pro" }, + "nebius/deepseek-ai/DeepSeek-V4-Pro-0813": { + "input_cost_per_token": 1.32e-06, + "litellm_provider": "nebius", + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/deepseek-ai%2FDeepSeek-V4-Pro-0813", + "supports_function_calling": true, + "supports_reasoning": true + }, "nebius/MiniMaxAI/MiniMax-M2.5": { "max_tokens": 196608, "max_input_tokens": 196608, @@ -37599,6 +37636,17 @@ "supports_reasoning": true, "source": "https://tokenfactory.nebius.com/models/catalog/text2text/zai-org%2FGLM-5.2" }, + "nebius/zai-org/GLM-5.3": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "nebius", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/zai-org%2FGLM-5.3", + "supports_function_calling": true, + "supports_reasoning": true + }, "nebius/zai-org/GLM-5.3-Flash": { "max_tokens": 1024000, "max_input_tokens": 1024000, @@ -44017,6 +44065,20 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "together_ai/deepseek-ai/DeepSeek-V4-Pro": { "deprecation_date": "2026-08-27", "cache_read_input_token_cost": 2e-07, @@ -61060,6 +61122,32 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/glm-5p3-fast": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/models/glm-5p3-flash": { "cache_read_input_token_cost": 3e-08, "input_cost_per_token": 1.5e-07, From 2a06ece58c1f830ad4ada60832bd0df5335b1cd3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 19:37:25 +0000 Subject: [PATCH 025/187] feat(models): mark Fireworks GLM 5.3 Fast router as supporting reasoning and response schema Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 ++++ model_prices_and_context_window.json | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0954f9360fa..67637285862 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -61132,6 +61132,8 @@ "output_cost_per_token": 6.6e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false }, @@ -61145,6 +61147,8 @@ "output_cost_per_token": 6.6e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false }, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0954f9360fa..67637285862 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -61132,6 +61132,8 @@ "output_cost_per_token": 6.6e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false }, @@ -61145,6 +61147,8 @@ "output_cost_per_token": 6.6e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false }, From 67382465860ecc96db6e1be3bdb68450e5223678 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:26:40 +0000 Subject: [PATCH 026/187] chore(models): credit absorbed Inception Mercury 2.5 cache pricing from #41016 Co-Authored-By: Nanduu24 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> From a478a46d0a6b306c999ed3cd5b367ff5b060ab03 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 21:17:12 -0700 Subject: [PATCH 027/187] Keep generated pricing parameters immutable --- .../pricing/test_price_precedence.py | 79 ++++++++++++++----- 1 file changed, 60 insertions(+), 19 deletions(-) diff --git a/tests/integration/pricing/test_price_precedence.py b/tests/integration/pricing/test_price_precedence.py index 1f724270a81..9f312c1e670 100644 --- a/tests/integration/pricing/test_price_precedence.py +++ b/tests/integration/pricing/test_price_precedence.py @@ -16,20 +16,36 @@ def test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic(ga @example(rates=(0, 0)) @example(rates=(1, 2)) @example(rates=("null", "null")) - @given(rates=st.one_of(st.sampled_from((("omitted", "omitted"), ("null", "null"))), st.tuples(st.integers(0, 25), st.integers(0, 25)))) + @given( + rates=st.one_of( + st.sampled_from((("omitted", "omitted"), ("null", "null"))), + st.tuples(st.integers(0, 25), st.integers(0, 25)), + ) + ) def check(rates: tuple[str | int, str | int]) -> None: - if rates[0] in ("omitted", "null"): - parameters: Final = {} if rates[0] == "omitted" else {"input_cost_per_token": None, "output_cost_per_token": None} - input_rate, output_rate = 0.00000015, 0.0000006 - else: - assert isinstance(rates[0], int) and isinstance(rates[1], int) - input_rate, output_rate = rates[0] / 1_000_000, rates[1] / 1_000_000 - parameters = {"input_cost_per_token": input_rate, "output_cost_per_token": output_rate} + defaults: Final = rates[0] in ("omitted", "null") + assert defaults or (isinstance(rates[0], int) and isinstance(rates[1], int)) + input_rate, output_rate = ( + (0.00000015, 0.0000006) if defaults else (float(rates[0]) / 1_000_000, float(rates[1]) / 1_000_000) + ) + parameters: Final = ( + {} + if rates[0] == "omitted" + else { + "input_cost_per_token": None if rates[0] == "null" else input_rate, + "output_cost_per_token": None if rates[0] == "null" else output_rate, + } + ) with gateway.scenario() as scenario: model: Final = scenario.model(**parameters) - response: Final = gateway.request("POST", "/v1/chat/completions", { - "model": model, "messages": [{"role": "user", "content": f"independent price {uuid.uuid4().hex}"}], - }) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"independent price {uuid.uuid4().hex}"}], + }, + ) assert response.status_code == 200, response.text assert response.json()["usage"] == {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40} expected: Final = 20 * input_rate + 20 * output_rate @@ -37,10 +53,14 @@ def test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic(ga assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6) else: assert response.headers.get("x-litellm-response-cost") in (None, "0", "0.0") - rows: Final = eventually(lambda: read_rows( - 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', - (response.json()["id"],), - ), lambda values: len(values) == 1, seconds=70) + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (response.json()["id"],), + ), + lambda values: len(values) == 1, + seconds=70, + ) assert rows[0]["prompt_tokens"] == 20 and rows[0]["completion_tokens"] == 20 assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6) metadata: Final = rows[0]["metadata"] @@ -56,18 +76,35 @@ def test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic(ga def test_same_upstream_aliases_keep_distinct_prices_after_reload(gateway: Gateway) -> None: for order in (("free", "paid"), ("paid", "free")): with gateway.scenario() as scenario: - rates: Final = {"free": {"input_cost_per_token": 0, "output_cost_per_token": 0}, "paid": {"input_cost_per_token": 0.001, "output_cost_per_token": 0.003}} + rates: Final = { + "free": {"input_cost_per_token": 0, "output_cost_per_token": 0}, + "paid": {"input_cost_per_token": 0.001, "output_cost_per_token": 0.003}, + } aliases: Final = {kind: scenario.model(**rates[kind]) for kind in order} for generation in range(2): for kind in order if generation == 0 else reversed(order): model: Final = aliases[kind] cost: Final = 0.08 if kind == "paid" else 0.0 - response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"alias price {model} {generation}"}]}) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"alias price {model} {generation}"}], + }, + ) assert response.status_code == 200, response.text assert response.json()["usage"]["total_tokens"] == 40 if cost: assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(cost) - rows: Final = eventually(lambda response=response: read_rows('SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response.json()["id"],)), lambda values: len(values) == 1, seconds=70) + rows: Final = eventually( + lambda response=response: read_rows( + 'SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (response.json()["id"],), + ), + lambda values: len(values) == 1, + seconds=70, + ) assert float(rows[0]["spend"]) == pytest.approx(cost) metadata: Final = rows[0]["metadata"] parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) @@ -77,5 +114,9 @@ def test_same_upstream_aliases_keep_distinct_prices_after_reload(gateway: Gatewa if generation == 0: entries: Final = gateway.get("/model/info")["data"] target: Final = next(entry for entry in entries if entry["model_name"] == aliases["paid"]) - changed: Final = gateway.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "price reload"}}) + changed: Final = gateway.request( + "PATCH", + f"/model/{target['model_info']['id']}/update", + {"model_info": {"description": "price reload"}}, + ) assert changed.status_code == 200, changed.text From dfac4e0a9f355123da4fe80421802d3a029ee711 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 21:24:20 -0700 Subject: [PATCH 028/187] Format database and accounting integration tests --- tests/integration/_support/process.py | 26 +++- .../database/test_partition_transactions.py | 27 ++-- .../test_reader_writer_regeneration.py | 87 ++++++++++--- .../database/test_transaction_atomicity.py | 92 +++++++++++--- .../pricing/test_price_precedence.py | 3 +- .../integration/spend/test_cache_and_quota.py | 116 ++++++++++++++---- 6 files changed, 280 insertions(+), 71 deletions(-) diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 7323152b0b6..0d66ecc9d90 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -62,10 +62,28 @@ def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str]) output.mkdir(parents=True, exist_ok=True) with (output / f"owned-proxy-{uuid.uuid4().hex}.log").open("w") as log: process: Final = subprocess.Popen( - [sys.executable, "-m", "integration._support.proxy", "--config", "tests/integration/proxy_config.yaml", - "--host", "127.0.0.1", "--port", str(port), "--num_workers", "1", "--telemetry", "False", - "--use_prisma_db_push", "--enforce_prisma_migration_check"], - cwd=root, env=environment, stdout=log, stderr=subprocess.STDOUT, start_new_session=True, + [ + sys.executable, + "-m", + "integration._support.proxy", + "--config", + "tests/integration/proxy_config.yaml", + "--host", + "127.0.0.1", + "--port", + str(port), + "--num_workers", + "1", + "--telemetry", + "False", + "--use_prisma_db_push", + "--enforce_prisma_migration_check", + ], + cwd=root, + env=environment, + stdout=log, + stderr=subprocess.STDOUT, + start_new_session=True, ) try: with httpx.Client(base_url=f"http://127.0.0.1:{port}", timeout=15, trust_env=False) as client: diff --git a/tests/integration/database/test_partition_transactions.py b/tests/integration/database/test_partition_transactions.py index 2759d8dcfb7..dbf54e6962b 100644 --- a/tests/integration/database/test_partition_transactions.py +++ b/tests/integration/database/test_partition_transactions.py @@ -21,17 +21,26 @@ class PartitionConnection: db: Prisma -@pytest.mark.covers("other.database.partitions.lock_wait_outlives_transaction_default", "other.database.partitions.repeat_preserves_rows") +@pytest.mark.covers( + "other.database.partitions.lock_wait_outlives_transaction_default", + "other.database.partitions.repeat_preserves_rows", +) async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> None: schema: Final = f"integration_{uuid.uuid4().hex}" url: Final = os.environ["DATABASE_URL"] parsed: Final = urlsplit(url) - scoped_url: Final = urlunsplit(parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema}))) + scoped_url: Final = urlunsplit( + parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema})) + ) parent: Final = sql.Identifier(schema, "LiteLLM_SpendLogs") with psycopg.connect(url, autocommit=True) as setup: setup.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema))) try: - setup.execute(sql.SQL('CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL) PARTITION BY RANGE ("startTime")').format(parent)) + setup.execute( + sql.SQL( + 'CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL) PARTITION BY RANGE ("startTime")' + ).format(parent) + ) database: Final = Prisma(datasource={"url": scoped_url}) await database.connect() try: @@ -39,12 +48,15 @@ async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> with psycopg.connect(url) as blocker: blocker.execute(sql.SQL("LOCK TABLE {} IN ACCESS SHARE MODE").format(parent)) blocker_pid: Final = blocker.info.backend_pid - operation: Final = asyncio.create_task(manager.ensure_partitions(PartitionConnection(database), lambda: 7000)) + operation: Final = asyncio.create_task( + manager.ensure_partitions(PartitionConnection(database), lambda: 7000) + ) wait_deadline: Final = time.monotonic() + 3 try: while True: witnesses: Final = read_rows( - "SELECT a.pid, extract(epoch FROM clock_timestamp()-a.query_start)::double precision AS age " + "SELECT a.pid, extract(epoch FROM " + "clock_timestamp()-a.query_start)::double precision AS age " "FROM pg_stat_activity a WHERE %s = ANY(pg_blocking_pids(a.pid)) " "AND a.wait_event_type = 'Lock' AND a.query LIKE 'CREATE TABLE IF NOT EXISTS%%'", (blocker_pid,), @@ -71,11 +83,12 @@ async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> catalog: Final = read_rows( "SELECT child.relname FROM pg_inherits i JOIN pg_class child ON child.oid=i.inhrelid " "JOIN pg_class parent ON parent.oid=i.inhparent JOIN pg_namespace n ON n.oid=parent.relnamespace " - "WHERE n.nspname=%s AND parent.relname='LiteLLM_SpendLogs'", (schema,), + "WHERE n.nspname=%s AND parent.relname='LiteLLM_SpendLogs'", + (schema,), ) assert catalog == [{"relname": ensured[0]}] now: Final = datetime.now(timezone.utc).replace(tzinfo=None) - setup.execute(sql.SQL('INSERT INTO {} VALUES (%s, %s)').format(parent), ("retained", now)) + setup.execute(sql.SQL("INSERT INTO {} VALUES (%s, %s)").format(parent), ("retained", now)) assert await manager.ensure_partitions(PartitionConnection(database), lambda: 7000) == ensured assert setup.execute(sql.SQL("SELECT request_id FROM {}").format(parent)).fetchall() == [("retained",)] finally: diff --git a/tests/integration/database/test_reader_writer_regeneration.py b/tests/integration/database/test_reader_writer_regeneration.py index 382dd14fbb0..2b4d93221ba 100644 --- a/tests/integration/database/test_reader_writer_regeneration.py +++ b/tests/integration/database/test_reader_writer_regeneration.py @@ -27,9 +27,15 @@ def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gatew role: Final = f"integration_reader_{uuid.uuid4().hex}" url: Final = os.environ["DATABASE_URL"] parsed: Final = urlsplit(url) - reader_url: Final = urlunsplit(parsed._replace(netloc=f"{role}:integration-reader-password@{parsed.hostname}:{parsed.port}")) + reader_url: Final = urlunsplit( + parsed._replace(netloc=f"{role}:integration-reader-password@{parsed.hostname}:{parsed.port}") + ) with psycopg.connect(url, autocommit=True) as admin: - admin.execute(sql.SQL("CREATE ROLE {} LOGIN PASSWORD 'integration-reader-password' NOSUPERUSER NOINHERIT").format(sql.Identifier(role))) + admin.execute( + sql.SQL("CREATE ROLE {} LOGIN PASSWORD 'integration-reader-password' NOSUPERUSER NOINHERIT").format( + sql.Identifier(role) + ) + ) try: admin.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(role))) admin.execute(sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA public TO {}").format(sql.Identifier(role))) @@ -39,7 +45,9 @@ def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gatew with pytest.raises(psycopg.errors.ReadOnlySqlTransaction): reader.execute('UPDATE "LiteLLM_VerificationToken" SET blocked = true WHERE false') with owned_proxy(gateway, tmp_path, {"DATABASE_URL_READ_REPLICA": reader_url}) as candidate: - assert read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,)), "Candidate reader was never connected" + assert read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,)), ( + "Candidate reader was never connected" + ) with gateway.scenario() as scenario: model: Final = scenario.model() outside: Final = scenario.model() @@ -48,12 +56,24 @@ def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gatew scenario.cleanups.callback(delete_if_present, gateway, old) scenario.cleanups.callback(delete_if_present, gateway, new) old_hash: Final = sha256(old.encode()).hexdigest() - before: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "no grant yet"}]}, key=old) - assert before.status_code == 403 and before.json()["error"]["type"] == "key_model_access_denied", before.text - response: Final = candidate.request("POST", "/v1/access_group", { - "access_group_name": f"integration-{uuid.uuid4().hex}", - "access_model_names": [model], "assigned_key_ids": [old_hash], - }) + before: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "no grant yet"}]}, + key=old, + ) + assert before.status_code == 403 and before.json()["error"]["type"] == "key_model_access_denied", ( + before.text + ) + response: Final = candidate.request( + "POST", + "/v1/access_group", + { + "access_group_name": f"integration-{uuid.uuid4().hex}", + "access_model_names": [model], + "assigned_key_ids": [old_hash], + }, + ) assert response.status_code == 201, response.text group: Final = string_value(response.json()["access_group_id"]) try: @@ -61,30 +81,57 @@ def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gatew blocker.execute('LOCK TABLE "LiteLLM_AccessGroupTable" IN ACCESS EXCLUSIVE MODE') pending: Final = executor.submit(candidate.request, "GET", f"/v1/access_group/{group}") try: - reached: Final = eventually(lambda: read_rows( - "SELECT usename FROM pg_stat_activity WHERE %s=ANY(pg_blocking_pids(pid)) " - "AND usename=%s AND query LIKE 'SELECT%%'", (blocker.info.backend_pid, role), - ), bool, seconds=3) + reached: Final = eventually( + lambda: read_rows( + "SELECT usename FROM pg_stat_activity WHERE %s=ANY(pg_blocking_pids(pid)) " + "AND usename=%s AND query LIKE 'SELECT%%'", + (blocker.info.backend_pid, role), + ), + bool, + seconds=3, + ) assert reached == [{"usename": role}] finally: blocker.rollback() selected: Final = pending.result(timeout=5) - assert selected.status_code == 200 and selected.json()["access_group_id"] == group, selected.text + assert selected.status_code == 200 and selected.json()["access_group_id"] == group, ( + selected.text + ) assert candidate.chat(model, key=old)["usage"]["total_tokens"] == 40 - regenerated: Final = candidate.post("/key/regenerate", {"key": old, "new_key": new, "grace_period": "0s"}) + regenerated: Final = candidate.post( + "/key/regenerate", {"key": old, "new_key": new, "grace_period": "0s"} + ) assert regenerated["key"] == new new_hash: Final = sha256(new.encode()).hexdigest() assert new != old - assert read_rows('SELECT assigned_key_ids FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)) == [{"assigned_key_ids": [new_hash]}] - assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s)', ([old_hash, new_hash],)) == [{"token": new_hash, "access_group_ids": [group]}] + assert read_rows( + 'SELECT assigned_key_ids FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,) + ) == [{"assigned_key_ids": [new_hash]}] + assert read_rows( + 'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s)', + ([old_hash, new_hash],), + ) == [{"token": new_hash, "access_group_ids": [group]}] assert candidate.chat(model, key=new)["usage"]["total_tokens"] == 40 assert candidate.chat(outside, key=new)["usage"]["total_tokens"] == 40 - denied: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rotated key"}]}, key=old) - assert denied.status_code == 401 and denied.json()["error"]["type"] == "token_not_found_in_db", denied.text + denied: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "rotated key"}]}, + key=old, + ) + assert ( + denied.status_code == 401 and denied.json()["error"]["type"] == "token_not_found_in_db" + ), denied.text finally: deleted: Final = gateway.request("DELETE", f"/v1/access_group/{group}") assert deleted.status_code == 204, deleted.text - assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)) == [] + assert ( + read_rows( + 'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', + (group,), + ) + == [] + ) finally: admin.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role))) admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role))) diff --git a/tests/integration/database/test_transaction_atomicity.py b/tests/integration/database/test_transaction_atomicity.py index 431d7982482..c150354d9a6 100644 --- a/tests/integration/database/test_transaction_atomicity.py +++ b/tests/integration/database/test_transaction_atomicity.py @@ -24,31 +24,80 @@ def test_access_group_second_key_constraint_failure_rolls_back_all_writes(gatewa witness: Final = constraint + "_seq" check_function: Final = constraint + "_check" body: Final = {"access_group_name": name, "access_model_names": [model], "assigned_key_ids": tokens} + def remove_partial_group() -> None: - for row in read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)): + for row in read_rows( + 'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,) + ): response: Final = gateway.request("DELETE", f"/v1/access_group/{row['access_group_id']}") assert response.status_code == 204, response.text - assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)) == [] + assert ( + read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)) + == [] + ) scenario.cleanups.callback(remove_partial_group) - before: Final = read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) + before: Final = read_rows( + 'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', + (tokens,), + ) with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection, ExitStack() as cleanup: connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(witness))) cleanup.callback(connection.execute, sql.SQL("DROP SEQUENCE {}").format(sql.Identifier(witness))) - connection.execute(sql.SQL("CREATE FUNCTION {}(text[]) RETURNS boolean LANGUAGE plpgsql AS $$ BEGIN IF cardinality($1)>0 THEN PERFORM nextval({}); RETURN false; END IF; RETURN true; END $$").format(sql.Identifier(check_function), sql.Literal(witness))) - cleanup.callback(connection.execute, sql.SQL("DROP FUNCTION {}(text[])").format(sql.Identifier(check_function))) - connection.execute(sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT {} CHECK (token <> {} OR {}(access_group_ids))').format(sql.Identifier(constraint), sql.Literal(tokens[1]), sql.Identifier(check_function))) - cleanup.callback(connection.execute, sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" DROP CONSTRAINT {}').format(sql.Identifier(constraint))) + connection.execute( + sql.SQL( + "CREATE FUNCTION {}(text[]) RETURNS boolean LANGUAGE plpgsql AS $$ BEGIN IF " + "cardinality($1)>0 THEN PERFORM nextval({}); RETURN false; END IF; RETURN true; END $$" + ).format(sql.Identifier(check_function), sql.Literal(witness)) + ) + cleanup.callback( + connection.execute, sql.SQL("DROP FUNCTION {}(text[])").format(sql.Identifier(check_function)) + ) + connection.execute( + sql.SQL( + 'ALTER TABLE "LiteLLM_VerificationToken" ADD ' + "CONSTRAINT {} CHECK (token <> {} OR {}(access_group_ids))" + ).format(sql.Identifier(constraint), sql.Literal(tokens[1]), sql.Identifier(check_function)) + ) + cleanup.callback( + connection.execute, + sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" DROP CONSTRAINT {}').format( + sql.Identifier(constraint) + ), + ) try: - assert connection.execute(sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness))).fetchone() == (False,) + assert connection.execute( + sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness)) + ).fetchone() == (False,) failed: Final = gateway.request("POST", "/v1/access_group", body) assert failed.status_code == 500, failed.text - assert connection.execute(sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness))).fetchone() == (True,) - assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)) == [] - assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) == before + assert connection.execute( + sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness)) + ).fetchone() == (True,) + assert ( + read_rows( + 'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,) + ) + == [] + ) + assert ( + read_rows( + "SELECT token, access_group_ids FROM " + '"LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', + (tokens,), + ) + == before + ) for key in keys: - denied: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rolled back grant"}]}, key=key) - assert denied.status_code == 403 and denied.json()["error"]["type"] == "key_model_access_denied", denied.text + denied: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "rolled back grant"}]}, + key=key, + ) + assert denied.status_code == 403 and denied.json()["error"]["type"] == "key_model_access_denied", ( + denied.text + ) finally: cleanup.close() created: Final = gateway.request("POST", "/v1/access_group", body) @@ -60,6 +109,17 @@ def test_access_group_second_key_constraint_failure_rolls_back_all_writes(gatewa finally: deleted: Final = gateway.request("DELETE", f"/v1/access_group/{identity}") assert deleted.status_code == 204, deleted.text - assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (identity,)) == [] - assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) == before - assert read_rows('SELECT conname FROM pg_constraint WHERE conname=%s', (constraint,)) == [] + assert ( + read_rows( + 'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (identity,) + ) + == [] + ) + assert ( + read_rows( + 'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', + (tokens,), + ) + == before + ) + assert read_rows("SELECT conname FROM pg_constraint WHERE conname=%s", (constraint,)) == [] diff --git a/tests/integration/pricing/test_price_precedence.py b/tests/integration/pricing/test_price_precedence.py index 9f312c1e670..0d73558d8a1 100644 --- a/tests/integration/pricing/test_price_precedence.py +++ b/tests/integration/pricing/test_price_precedence.py @@ -55,7 +55,8 @@ def test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic(ga assert response.headers.get("x-litellm-response-cost") in (None, "0", "0.0") rows: Final = eventually( lambda: read_rows( - 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + "SELECT spend, metadata, prompt_tokens, " + 'completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response.json()["id"],), ), lambda values: len(values) == 1, diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 05bf100f091..114aabbae33 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -22,7 +22,9 @@ def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gate self.resources = ExitStack() try: self.scenario = self.resources.enter_context(gateway.scenario()) - self.upstream = self.resources.enter_context(httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False)) + self.upstream = self.resources.enter_context( + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) + ) self.model = self.scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) self.key = self.scenario.key(models=[self.model]) self.prefix = uuid.uuid4().hex @@ -46,13 +48,22 @@ def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gate def perform_request(self, marker: int) -> None: self.upstream.get("/__observations").raise_for_status() - response: Final = gateway.request("POST", "/v1/chat/completions", { - "model": self.model, "messages": [{"role": "user", "content": f"{self.prefix}-{marker}"}], - }, key=self.key) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": self.model, + "messages": [{"role": "user", "content": f"{self.prefix}-{marker}"}], + }, + key=self.key, + ) assert response.status_code == 200, response.text self.requests += 1 body: Final = response.json() - assert body["choices"][0]["message"]["content"] == "Hello! This is a mock response from the fake OpenAI endpoint." + assert ( + body["choices"][0]["message"]["content"] + == "Hello! This is a mock response from the fake OpenAI endpoint." + ) assert body["usage"]["total_tokens"] == 40 observed: Final = self.upstream.get("/__observations").json()["requests"] expected_calls: Final = 0 if marker in self.seen else 1 @@ -70,10 +81,15 @@ def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gate def teardown(self) -> None: try: if self.requests and not self.failed: - rows: Final = eventually(lambda: read_rows( - 'SELECT request_id, spend, cache_hit, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', - (sha256(self.key.encode()).hexdigest(),), - ), lambda values: len(values) == self.requests, seconds=70) + rows: Final = eventually( + lambda: read_rows( + "SELECT request_id, spend, cache_hit, prompt_tokens, " + 'completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(self.key.encode()).hexdigest(),), + ), + lambda values: len(values) == self.requests, + seconds=70, + ) assert len({row["request_id"] for row in rows}) == self.requests assert sum(float(row["spend"]) for row in rows) == pytest.approx(self.paid * 0.06) assert sum(row["cache_hit"] == "True" for row in rows) == self.requests - self.paid @@ -81,7 +97,10 @@ def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gate assert row["prompt_tokens"] == 20 and row["completion_tokens"] == 20 if row["cache_hit"] == "True": assert float(row["spend"]) == 0 and "_cache_hit" in row["request_id"] - assert any(row["request_id"].startswith(identity + "_cache_hit") for identity in self.identities.values()) + assert any( + row["request_id"].startswith(identity + "_cache_hit") + for identity in self.identities.values() + ) else: assert row["request_id"] in self.identities.values() assert float(row["spend"]) == pytest.approx(0.06) @@ -95,7 +114,10 @@ def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gate @pytest.mark.covers("quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge") def test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows(gateway: Gateway) -> None: - with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) key: Final = scenario.key(models=[model]) prompt: Final = f"repeated cache {uuid.uuid4().hex}" @@ -104,9 +126,19 @@ def test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows assert len(upstream.get("/__observations").json()["requests"]) == 1 assert len({result["id"] for result in results}) == 1 for result in results: - assert result["choices"][0]["message"]["content"] == "Hello! This is a mock response from the fake OpenAI endpoint." + assert ( + result["choices"][0]["message"]["content"] + == "Hello! This is a mock response from the fake OpenAI endpoint." + ) assert result["usage"]["total_tokens"] == 40 - rows: Final = eventually(lambda: read_rows('SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)), lambda values: len(values) == 3, seconds=70) + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ), + lambda values: len(values) == 3, + seconds=70, + ) assert len({row["request_id"] for row in rows}) == 3 assert sorted(float(row["spend"]) for row in rows) == [0, 0, 0.06] for row in rows: @@ -119,39 +151,74 @@ def test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows @pytest.mark.covers("quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores") def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gateway: Gateway) -> None: - with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) key: Final = scenario.key(models=[model], max_budget=0.06) control: Final = scenario.key(models=[model]) first: Final = gateway.chat(model, key=key, text=f"budget {uuid.uuid4().hex}") assert first["usage"]["total_tokens"] == 40 digest: Final = sha256(key.encode()).hexdigest() - spent: Final = eventually(lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, seconds=70) + spent: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, + seconds=70, + ) assert float(spent[0]["spend"]) == pytest.approx(0.06) upstream.get("/__observations").raise_for_status() - denied: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"over budget {uuid.uuid4().hex}"}]}, key=key) + denied: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"over budget {uuid.uuid4().hex}"}]}, + key=key, + ) assert denied.status_code == 429 and denied.json()["error"]["type"] == "budget_exceeded", denied.text assert upstream.get("/__observations").json()["requests"] == [] assert gateway.chat(model, key=control, text=f"control {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 gateway.post("/key/update", {"key": key, "spend": 0}) - assert read_rows('SELECT spend, max_budget FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [{"spend": 0.0, "max_budget": 0.06}] + assert read_rows('SELECT spend, max_budget FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [ + {"spend": 0.0, "max_budget": 0.06} + ] assert gateway.chat(model, key=key, text=f"reset {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 - eventually(lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, seconds=70) + eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, + seconds=70, + ) upstream.get("/__observations").raise_for_status() - denied_again: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"boundary again {uuid.uuid4().hex}"}]}, key=key) - assert denied_again.status_code == 429 and denied_again.json()["error"]["type"] == "budget_exceeded", denied_again.text + denied_again: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"boundary again {uuid.uuid4().hex}"}]}, + key=key, + ) + assert denied_again.status_code == 429 and denied_again.json()["error"]["type"] == "budget_exceeded", ( + denied_again.text + ) assert upstream.get("/__observations").json()["requests"] == [] @pytest.mark.covers("quota_management.response_cache.system_messages_partition_cache_identity") def test_different_system_messages_do_not_share_a_cached_response(gateway: Gateway) -> None: - with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): model: Final = scenario.model() prompt: Final = uuid.uuid4().hex identities: dict[str, str] = {} for system, expected_calls in (("first policy", 1), ("second policy", 1), ("first policy", 0)): upstream.get("/__observations").raise_for_status() - response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "system", "content": system}, {"role": "user", "content": prompt}]}) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "system", "content": system}, {"role": "user", "content": prompt}], + }, + ) assert response.status_code == 200 and response.json()["usage"]["total_tokens"] == 40, response.text calls: Final = upstream.get("/__observations").json()["requests"] assert len(calls) == expected_calls @@ -161,4 +228,7 @@ def test_different_system_messages_do_not_share_a_cached_response(gateway: Gatew assert response.json()["id"] not in identities.values() identities = {**identities, system: response.json()["id"]} if calls: - assert calls[0]["body"]["messages"] == [{"role": "system", "content": system}, {"role": "user", "content": prompt}] + assert calls[0]["body"]["messages"] == [ + {"role": "system", "content": system}, + {"role": "user", "content": prompt}, + ] From e7454e52e6b15c9d7cce043abb84d43eecb1f0d7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:09:05 +0000 Subject: [PATCH 029/187] fix(gemini): map minimal thinking to low for Gemini 3.7 and 3.8 Flash Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../vertex_and_google_ai_studio_gemini.py | 35 ++++----- ...test_vertex_and_google_ai_studio_gemini.py | 78 +++++++++++++++++++ 2 files changed, 95 insertions(+), 18 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d113b2b4f6b..d719d53e19f 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -110,6 +110,7 @@ else: SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "none", "disable") +GEMINI_FLASH_MODELS_WITHOUT_MINIMAL_THINKING: Final = ("gemini-3.7-flash", "gemini-3.8-flash") def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsError: @@ -860,6 +861,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): else: raise _unsupported_reasoning_effort(reasoning_effort) + @staticmethod + def _supports_minimal_thinking_level(model: str) -> bool: + lowered: Final = model.lower() + is_gemini3flash: Final = "gemini-3" in lowered and "flash" in lowered + return is_gemini3flash and not any(m in lowered for m in GEMINI_FLASH_MODELS_WITHOUT_MINIMAL_THINKING) + @staticmethod def _map_reasoning_effort_to_thinking_level( reasoning_effort: str, @@ -874,13 +881,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): Returns: GeminiThinkingConfig with thinkingLevel and includeThoughts """ - # Check if this is gemini-3-flash which supports MINIMAL thinking level - # Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview, - # gemini-3.5-flash, and any future 3.x-flash variants. is_gemini3flash: Final = model and ("flash" in model.lower() and "gemini-3" in model.lower()) + supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model) is_gemini31pro: Final = model and ("gemini-3.1-pro-preview" in model.lower()) if reasoning_effort == "minimal": - if is_gemini3flash: + if supports_minimal: return {"thinkingLevel": "minimal", "includeThoughts": True} else: return {"thinkingLevel": "low", "includeThoughts": True} @@ -893,18 +898,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return {"thinkingLevel": "high", "includeThoughts": True} elif reasoning_effort == "high": return {"thinkingLevel": "high", "includeThoughts": True} - elif reasoning_effort == "disable": - # Gemini 3 cannot fully disable thinking, so we use "minimal" for gemini-3-flash-preview, "low" for others - if is_gemini3flash: - return {"thinkingLevel": "minimal", "includeThoughts": False} - else: - return {"thinkingLevel": "low", "includeThoughts": False} - elif reasoning_effort == "none": - # For gemini-3-flash-preview, use "minimal" instead of "low" - if is_gemini3flash: - return {"thinkingLevel": "minimal", "includeThoughts": False} - else: - return {"thinkingLevel": "low", "includeThoughts": False} + elif reasoning_effort in ("disable", "none"): + return { + "thinkingLevel": "minimal" if supports_minimal else "low", + "includeThoughts": False, + } else: raise _unsupported_reasoning_effort(reasoning_effort) @@ -971,8 +969,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): params["includeThoughts"] = True # Follow provider defaults unless explicitly opted into legacy behavior. if litellm.enable_gemini_default_thinking_level_low is True: - is_gemini3flash: Final = "gemini-3" in model.lower() and "flash" in model.lower() - params["thinkingLevel"] = "minimal" if is_gemini3flash else "low" + params["thinkingLevel"] = ( + "minimal" if VertexGeminiConfig._supports_minimal_thinking_level(model) else "low" + ) else: # Thinking disabled params["includeThoughts"] = False diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 101f6e6fa5d..a1c31689d09 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2678,6 +2678,84 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): assert result["thinkingConfig"]["includeThoughts"] is False +@pytest.mark.parametrize( + "model", + [ + "gemini-3.7-flash", + "vertex_ai/gemini-3.8-flash", + "gemini-3.8-flash-preview", + ], +) +@pytest.mark.parametrize( + ("reasoning_effort", "include_thoughts"), + [("minimal", True), ("none", False), ("disable", False)], +) +def test_gemini_37_38_flash_floor_minimal_thinking_level( + model, reasoning_effort, include_thoughts +): + result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( + reasoning_effort, model + ) + + assert result["thinkingLevel"] == "low" + assert result["includeThoughts"] is include_thoughts + + +@pytest.mark.parametrize( + ("model", "reasoning_effort", "expected_level", "include_thoughts"), + [ + ("gemini-3-flash-preview", "minimal", "minimal", True), + ("gemini-3-flash-preview", "none", "minimal", False), + ("gemini-3-flash-preview", "disable", "minimal", False), + ("gemini-3.6-flash", "minimal", "minimal", True), + ("gemini-3.6-flash", "none", "minimal", False), + ("gemini-3.6-flash", "disable", "minimal", False), + ("gemini-3.5-flash", "minimal", "minimal", True), + ("gemini-3.5-flash", "none", "minimal", False), + ("gemini-3.5-flash", "disable", "minimal", False), + ("gemini-3.8-flash", "medium", "medium", True), + ], +) +def test_gemini_flash_minimal_thinking_support( + model, reasoning_effort, expected_level, include_thoughts +): + result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( + reasoning_effort, model + ) + + assert result["thinkingLevel"] == expected_level + assert result["includeThoughts"] is include_thoughts + + +def test_gemini_38_flash_feature_flag_uses_low_thinking_level(monkeypatch): + monkeypatch.setattr(litellm, "enable_gemini_default_thinking_level_low", True) + thinking_param = {"type": "enabled", "budget_tokens": 1024} + + result_38 = VertexGeminiConfig._map_thinking_param( + thinking_param, model="gemini-3.8-flash" + ) + result_36 = VertexGeminiConfig._map_thinking_param( + thinking_param, model="gemini-3.6-flash" + ) + + assert result_38["thinkingLevel"] == "low" + assert result_36["thinkingLevel"] == "minimal" + + +def test_gemini_38_flash_public_reasoning_effort_none_uses_low(): + result = VertexGeminiConfig().map_openai_params( + non_default_params={"reasoning_effort": "none"}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + + assert result["thinkingConfig"] == { + "thinkingLevel": "low", + "includeThoughts": False, + } + + def test_reasoning_effort_dict_format_gemini_3(): """ Test that reasoning_effort works when passed as dict format from OpenAI Agents SDK. From 631ba7f9e3d1d8dc8f3bc3616ad33723e307d305 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:36:11 +0000 Subject: [PATCH 030/187] fix(models): dedupe merged keys, price gemini *-latest aliases at live targets, add Nova cache read prices and Fireworks deprecation dates Absorbs #41148 and #41152 into the rolling registry PR. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 338 +++++++++++------- model_prices_and_context_window.json | 338 +++++++++++------- ..._cross_region_inference_profile_mapping.py | 54 ++- .../llms/gemini/test_cost_calculator.py | 42 ++- 4 files changed, 495 insertions(+), 277 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 39370862516..d3c3cb9079b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -376,7 +376,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -558,7 +559,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "amazon.nova-pro-v1:0": { "input_cost_per_token": 8e-07, @@ -573,7 +575,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -22470,7 +22473,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "deprecation_date": "2026-08-27" }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, @@ -22857,7 +22861,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "deprecation_date": "2026-08-27" }, "fireworks_ai/accounts/fireworks/models/minimax-m3": { "cache_read_input_token_cost": 6e-08, @@ -22963,7 +22968,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "deprecation_date": "2026-08-27" }, "fireworks_ai/glm-4p7": { "cache_read_input_token_cost": 3e-07, @@ -23181,7 +23187,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "deprecation_date": "2026-08-27" }, "fireworks_ai/minimax-m3": { "cache_read_input_token_cost": 6e-08, @@ -26435,13 +26442,71 @@ "supports_image_size": false }, "gemini/gemini-flash-latest": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "rpm": 15, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true, + "tpm": 250000, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.35e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "output_cost_per_token_priority": 6.75e-06, + "prompt_cache_min_tokens": 4096, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" + }, + "gemini/gemini-flash-lite-latest": { "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, @@ -26475,58 +26540,23 @@ "supports_web_search": true, "tpm": 250000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 - }, - "gemini/gemini-flash-lite-latest": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, - "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "rpm": 15, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000, - "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 - }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini/gemini-2.5-flash-lite-preview-06-17": { "deprecation_date": "2025-11-18", @@ -44717,7 +44747,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "us.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.5e-08, @@ -44730,7 +44761,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -44743,9 +44775,10 @@ "output_cost_per_token": 1.25e-05, "supports_function_calling": true, "supports_pdf_input": true, - "supports_prompt_caching": false, + "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 6.25e-07 }, "us.amazon.nova-pro-v1:0": { "input_cost_per_token": 8e-07, @@ -44760,7 +44793,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, @@ -55684,18 +55718,17 @@ ] }, "gemini-flash-latest": { - "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -55724,25 +55757,37 @@ "supports_web_search": true, "tpm": 8000000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.35e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "output_cost_per_token_priority": 6.75e-06, + "prompt_cache_min_tokens": 4096, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini-flash-lite-latest": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, "rpm": 15, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-lite", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -55771,29 +55816,42 @@ "supports_web_search": true, "tpm": 250000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini-pro-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "rpm": 2000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -55817,29 +55875,42 @@ "supports_web_search": true, "tpm": 800000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096, + "supports_native_streaming": true, + "supports_url_context": true, + "web_search_billing_unit": "per_query" }, "gemini/gemini-pro-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "rpm": 2000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -55863,11 +55934,23 @@ "supports_web_search": true, "tpm": 800000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096, + "supports_native_streaming": true, + "supports_url_context": true, + "web_search_billing_unit": "per_query" }, "gemini-exp-1206": { "cache_read_input_token_cost": 3e-08, @@ -58471,7 +58554,8 @@ "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_response_schema": true }, "fireworks_ai/accounts/fireworks/models/kimi-k3": { "cache_read_input_token_cost": 3e-07, @@ -65873,14 +65957,6 @@ "supports_response_schema": true, "supports_vision": false }, - "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { - "cache_read_input_token_cost": 3.9e-07, - "input_cost_per_token": 2.1e-06, - "litellm_provider": "fireworks_ai", - "mode": "chat", - "output_cost_per_token": 6.6e-06, - "source": "https://api.fireworks.ai/v1/serverless/models" - }, "together_ai/arcee-ai/trinity-mini": { "input_cost_per_token": 4.5e-08, "litellm_provider": "together_ai", @@ -65916,14 +65992,6 @@ "output_cost_per_token": 1.6e-06, "source": "https://api.together.ai/v1/models" }, - "together_ai/deepseek-ai/DeepSeek-V4.1-Flash": { - "cache_read_input_token_cost": 6e-09, - "input_cost_per_token": 3e-07, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1.2e-06, - "source": "https://api.together.ai/v1/models" - }, "vertex_ai/gemini-2.5-flash-native-audio": { "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 39370862516..d3c3cb9079b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -376,7 +376,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -558,7 +559,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "amazon.nova-pro-v1:0": { "input_cost_per_token": 8e-07, @@ -573,7 +575,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -22470,7 +22473,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "deprecation_date": "2026-08-27" }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, @@ -22857,7 +22861,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "deprecation_date": "2026-08-27" }, "fireworks_ai/accounts/fireworks/models/minimax-m3": { "cache_read_input_token_cost": 6e-08, @@ -22963,7 +22968,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "deprecation_date": "2026-08-27" }, "fireworks_ai/glm-4p7": { "cache_read_input_token_cost": 3e-07, @@ -23181,7 +23187,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "deprecation_date": "2026-08-27" }, "fireworks_ai/minimax-m3": { "cache_read_input_token_cost": 6e-08, @@ -26435,13 +26442,71 @@ "supports_image_size": false }, "gemini/gemini-flash-latest": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "rpm": 15, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true, + "tpm": 250000, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.35e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "output_cost_per_token_priority": 6.75e-06, + "prompt_cache_min_tokens": 4096, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" + }, + "gemini/gemini-flash-lite-latest": { "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, @@ -26475,58 +26540,23 @@ "supports_web_search": true, "tpm": 250000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 - }, - "gemini/gemini-flash-lite-latest": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, - "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "rpm": 15, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000, - "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 - }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini/gemini-2.5-flash-lite-preview-06-17": { "deprecation_date": "2025-11-18", @@ -44717,7 +44747,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "us.amazon.nova-micro-v1:0": { "input_cost_per_token": 3.5e-08, @@ -44730,7 +44761,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -44743,9 +44775,10 @@ "output_cost_per_token": 1.25e-05, "supports_function_calling": true, "supports_pdf_input": true, - "supports_prompt_caching": false, + "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 6.25e-07 }, "us.amazon.nova-pro-v1:0": { "input_cost_per_token": 8e-07, @@ -44760,7 +44793,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, @@ -55684,18 +55718,17 @@ ] }, "gemini-flash-latest": { - "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -55724,25 +55757,37 @@ "supports_web_search": true, "tpm": 8000000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.35e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "output_cost_per_token_priority": 6.75e-06, + "prompt_cache_min_tokens": 4096, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini-flash-lite-latest": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, "rpm": 15, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-lite", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -55771,29 +55816,42 @@ "supports_web_search": true, "tpm": 250000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini-pro-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "rpm": 2000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -55817,29 +55875,42 @@ "supports_web_search": true, "tpm": 800000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096, + "supports_native_streaming": true, + "supports_url_context": true, + "web_search_billing_unit": "per_query" }, "gemini/gemini-pro-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "rpm": 2000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -55863,11 +55934,23 @@ "supports_web_search": true, "tpm": 800000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096, + "supports_native_streaming": true, + "supports_url_context": true, + "web_search_billing_unit": "per_query" }, "gemini-exp-1206": { "cache_read_input_token_cost": 3e-08, @@ -58471,7 +58554,8 @@ "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_response_schema": true }, "fireworks_ai/accounts/fireworks/models/kimi-k3": { "cache_read_input_token_cost": 3e-07, @@ -65873,14 +65957,6 @@ "supports_response_schema": true, "supports_vision": false }, - "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { - "cache_read_input_token_cost": 3.9e-07, - "input_cost_per_token": 2.1e-06, - "litellm_provider": "fireworks_ai", - "mode": "chat", - "output_cost_per_token": 6.6e-06, - "source": "https://api.fireworks.ai/v1/serverless/models" - }, "together_ai/arcee-ai/trinity-mini": { "input_cost_per_token": 4.5e-08, "litellm_provider": "together_ai", @@ -65916,14 +65992,6 @@ "output_cost_per_token": 1.6e-06, "source": "https://api.together.ai/v1/models" }, - "together_ai/deepseek-ai/DeepSeek-V4.1-Flash": { - "cache_read_input_token_cost": 6e-09, - "input_cost_per_token": 3e-07, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1.2e-06, - "source": "https://api.together.ai/v1/models" - }, "vertex_ai/gemini-2.5-flash-native-audio": { "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index fda3c8ceb8f..b567a8669b4 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -4,12 +4,10 @@ from typing import NamedTuple import pytest - import litellm +from litellm.cost_calculator import completion_cost from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig from litellm.llms.bedrock.common_utils import BedrockModelInfo -from litellm.utils import _get_model_info_helper -from litellm.cost_calculator import completion_cost from litellm.types.utils import ( Choices, Message, @@ -17,6 +15,7 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, Usage, ) +from litellm.utils import _get_model_info_helper @pytest.fixture @@ -228,3 +227,52 @@ def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(prof assert "reasoning_effort" in supported assert "thinking" not in supported assert "output_config" not in supported + + +@pytest.mark.parametrize( + "model,expected_cache_read", + [ + ("amazon.nova-lite-v1:0", 1.5e-8), + ("us.amazon.nova-lite-v1:0", 1.5e-8), + ("amazon.nova-micro-v1:0", 8.75e-9), + ("us.amazon.nova-micro-v1:0", 8.75e-9), + ("amazon.nova-pro-v1:0", 2e-7), + ("us.amazon.nova-pro-v1:0", 2e-7), + ("us.amazon.nova-premier-v1:0", 6.25e-7), + ], +) +def test_bedrock_nova_cache_read_prices( + model, expected_cache_read, local_model_cost_map +): + model_info = litellm.model_cost[model] + usage = Usage( + prompt_tokens=1_000, + completion_tokens=100, + total_tokens=1_100, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=400), + ) + response = _bedrock_response(model, usage) + + cost = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + ) + expected_cost = ( + 600 * model_info["input_cost_per_token"] + + 400 * expected_cache_read + + 100 * model_info["output_cost_per_token"] + ) + assert cost == pytest.approx(expected_cost) + + uncached_usage = Usage( + prompt_tokens=1_000, + completion_tokens=100, + total_tokens=1_100, + ) + uncached_cost = completion_cost( + completion_response=_bedrock_response(model, uncached_usage), + model=model, + custom_llm_provider="bedrock", + ) + assert cost < uncached_cost diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index 6d547b0dc55..85d62030c00 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -1,8 +1,7 @@ -import os - import pytest import litellm +from litellm.cost_calculator import completion_cost from litellm.llms.gemini.cost_calculator import ( cost_per_google_maps_grounding_request, cost_per_web_search_request, @@ -18,6 +17,7 @@ from litellm.types.utils import ( ImageResponse, ImageUsage, ImageUsageInputTokensDetails, + ModelResponse, PromptTokensDetailsWrapper, Usage, ) @@ -455,8 +455,8 @@ def test_map_traffic_type_to_service_tier( @pytest.mark.parametrize( "model,custom_llm_provider,expected_cache_read_cost", [ - ("gemini/gemini-flash-latest", "gemini", 3e-08), - ("gemini/gemini-flash-lite-latest", "gemini", 1e-08), + ("gemini/gemini-flash-latest", "gemini", 7.5e-08), + ("gemini/gemini-flash-lite-latest", "gemini", 3e-08), ("gemini/gemini-2.5-flash-preview-09-2025", "gemini", 3e-08), ("gemini/gemini-2.5-flash-lite-preview-06-17", "gemini", 1e-08), ("vertex_ai/gemini-2.5-flash-preview-09-2025", "vertex_ai", 3e-08), @@ -479,6 +479,40 @@ def test_flash_alias_cache_read_is_ten_percent_of_input( ) +@pytest.mark.parametrize( + "alias,target", + [ + ("gemini/gemini-flash-latest", "gemini/gemini-3.8-flash"), + ("gemini/gemini-flash-lite-latest", "gemini/gemini-3.5-flash-lite"), + ("gemini/gemini-pro-latest", "gemini/gemini-3.1-pro-preview"), + ], +) +def test_latest_aliases_cost_the_same_as_their_current_target( + monkeypatch, alias, target +): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + usage = Usage( + prompt_tokens=1_000, + completion_tokens=500, + total_tokens=1_500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=400), + ) + + def cost_of(model: str) -> float: + return completion_cost( + completion_response=ModelResponse(model=model, usage=usage), + model=model, + custom_llm_provider="gemini", + ) + + alias_cost = cost_of(alias) + target_cost = cost_of(target) + assert alias_cost == pytest.approx(target_cost) + assert alias_cost > 0 + + @pytest.mark.parametrize( "prefixed,bare", [ From 3000c00c502d4d662ba369b2fa3bd628f73b9cc8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 14:03:24 +0000 Subject: [PATCH 031/187] fix(models): add supports_response_schema to fireworks deepseek-v4-flash-vision-exp short alias Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 3 ++- model_prices_and_context_window.json | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d3c3cb9079b..e2c9a3edb87 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -58631,7 +58631,8 @@ "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_response_schema": true }, "fireworks_ai/glm-5p2-fast": { "cache_read_input_token_cost": 2.1e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d3c3cb9079b..e2c9a3edb87 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -58631,7 +58631,8 @@ "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_response_schema": true }, "fireworks_ai/glm-5p2-fast": { "cache_read_input_token_cost": 2.1e-07, From 93f27d6082f235bb7c264abad0006038d16b321d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:15:46 +0000 Subject: [PATCH 032/187] fix(registry): add gemini 3.8 live, azure gpt-5.5/luna snapshots, doubao seed 2.1, fix together v4.1 flash context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 425 +++++++++++++++++- model_prices_and_context_window.json | 425 +++++++++++++++++- 2 files changed, 844 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e2c9a3edb87..ce531ba1fe1 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3781,6 +3781,53 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure_ai/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure_ai/gpt-5.4": { "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, @@ -7380,6 +7427,62 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-luna-2026-07-09": { + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_priority": 5e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "cache_read_input_token_cost_priority": 4e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "input_cost_per_token_priority": 4e-07, + "input_cost_per_token_above_272k_tokens_priority": 8e-07, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "output_cost_per_token_priority": 2.4e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, @@ -8218,6 +8321,58 @@ "supports_web_search": true, "deprecation_date": "2027-10-26" }, + "azure/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/us/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, @@ -8263,6 +8418,55 @@ "supports_web_search": true, "deprecation_date": "2027-10-26" }, + "azure/us/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "cache_read_input_token_cost_priority": 1.38e-06, + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "input_cost_per_token_priority": 1.375e-05, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "output_cost_per_token_priority": 8.25e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/eu/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, @@ -8308,6 +8512,55 @@ "supports_web_search": true, "deprecation_date": "2027-10-26" }, + "azure/eu/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "cache_read_input_token_cost_priority": 1.38e-06, + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "input_cost_per_token_priority": 1.375e-05, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "output_cost_per_token_priority": 8.25e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, @@ -44436,11 +44689,11 @@ "cache_read_input_token_cost": 6e-09, "input_cost_per_token": 3e-07, "litellm_provider": "together_ai", - "max_input_tokens": 1000000, - "max_tokens": 1000000, + "max_input_tokens": 1048576, + "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.xyz/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -55573,6 +55826,73 @@ "supports_web_search": true, "gemini_audio_only_live": true }, + "gemini-3.8-live": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true + }, + "gemini-3.8-live-extended-thinking": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true, + "supports_reasoning": true + }, "gemini/gemini-2.5-flash-native-audio-latest": { "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, @@ -55692,6 +56012,73 @@ "rpm": 10, "gemini_audio_only_live": true }, + "gemini/gemini-3.8-live": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true + }, + "gemini/gemini-3.8-live-extended-thinking": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true, + "supports_reasoning": true + }, "gemini/gemini-3.1-flash-tts-preview": { "input_cost_per_token": 1e-06, "litellm_provider": "gemini", @@ -56850,6 +57237,38 @@ } ] }, + "volcengine/doubao-seed-2-1-pro-260628": { + "cache_read_input_token_cost": 1.725e-07, + "input_cost_per_token": 8.625e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4.3125e-06, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, + "volcengine/doubao-seed-2-1-turbo-260628": { + "cache_read_input_token_cost": 8.625e-08, + "input_cost_per_token": 4.3125e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 2.15625e-06, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, "volcengine/doubao-seed-2-0-lite-260215": { "litellm_provider": "volcengine", "max_input_tokens": 256000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e2c9a3edb87..ce531ba1fe1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3781,6 +3781,53 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure_ai/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure_ai/gpt-5.4": { "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, @@ -7380,6 +7427,62 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-luna-2026-07-09": { + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_priority": 5e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "cache_read_input_token_cost_priority": 4e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "input_cost_per_token_priority": 4e-07, + "input_cost_per_token_above_272k_tokens_priority": 8e-07, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "output_cost_per_token_priority": 2.4e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, @@ -8218,6 +8321,58 @@ "supports_web_search": true, "deprecation_date": "2027-10-26" }, + "azure/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/us/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, @@ -8263,6 +8418,55 @@ "supports_web_search": true, "deprecation_date": "2027-10-26" }, + "azure/us/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "cache_read_input_token_cost_priority": 1.38e-06, + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "input_cost_per_token_priority": 1.375e-05, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "output_cost_per_token_priority": 8.25e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/eu/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, @@ -8308,6 +8512,55 @@ "supports_web_search": true, "deprecation_date": "2027-10-26" }, + "azure/eu/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "cache_read_input_token_cost_priority": 1.38e-06, + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "input_cost_per_token_priority": 1.375e-05, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "output_cost_per_token_priority": 8.25e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, @@ -44436,11 +44689,11 @@ "cache_read_input_token_cost": 6e-09, "input_cost_per_token": 3e-07, "litellm_provider": "together_ai", - "max_input_tokens": 1000000, - "max_tokens": 1000000, + "max_input_tokens": 1048576, + "max_tokens": 1048576, "mode": "chat", "output_cost_per_token": 1.2e-06, - "source": "https://docs.together.ai/docs/serverless-models", + "source": "https://api.together.xyz/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, @@ -55573,6 +55826,73 @@ "supports_web_search": true, "gemini_audio_only_live": true }, + "gemini-3.8-live": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true + }, + "gemini-3.8-live-extended-thinking": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true, + "supports_reasoning": true + }, "gemini/gemini-2.5-flash-native-audio-latest": { "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, @@ -55692,6 +56012,73 @@ "rpm": 10, "gemini_audio_only_live": true }, + "gemini/gemini-3.8-live": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true + }, + "gemini/gemini-3.8-live-extended-thinking": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true, + "supports_reasoning": true + }, "gemini/gemini-3.1-flash-tts-preview": { "input_cost_per_token": 1e-06, "litellm_provider": "gemini", @@ -56850,6 +57237,38 @@ } ] }, + "volcengine/doubao-seed-2-1-pro-260628": { + "cache_read_input_token_cost": 1.725e-07, + "input_cost_per_token": 8.625e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4.3125e-06, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, + "volcengine/doubao-seed-2-1-turbo-260628": { + "cache_read_input_token_cost": 8.625e-08, + "input_cost_per_token": 4.3125e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 2.15625e-06, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, "volcengine/doubao-seed-2-0-lite-260215": { "litellm_provider": "volcengine", "max_input_tokens": 256000, From 11211cbf29087a5eb474dd708e04a9910494ddd3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:18:29 +0000 Subject: [PATCH 033/187] fix(registry): drop gemini/gemini-3.8-live entries pending published rate limits The gemini/ realtime test requires tpm and rpm, and Google publishes the Live model limits only behind the AI Studio login, so the direct-API keys cannot be sourced yet. The Vertex keys stay. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 67 ------------------- model_prices_and_context_window.json | 67 ------------------- 2 files changed, 134 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ce531ba1fe1..5631220144e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -56012,73 +56012,6 @@ "rpm": 10, "gemini_audio_only_live": true }, - "gemini/gemini-3.8-live": { - "input_cost_per_audio_token": 3e-06, - "input_cost_per_image_token": 1e-06, - "input_cost_per_token": 7.5e-07, - "input_cost_per_video_per_second": 3.3333333333333335e-05, - "litellm_provider": "gemini", - "max_input_tokens": 131072, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "realtime", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 4.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/realtime" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_vision": true, - "supports_web_search": true, - "gemini_audio_only_live": true - }, - "gemini/gemini-3.8-live-extended-thinking": { - "input_cost_per_audio_token": 3e-06, - "input_cost_per_image_token": 1e-06, - "input_cost_per_token": 7.5e-07, - "input_cost_per_video_per_second": 3.3333333333333335e-05, - "litellm_provider": "gemini", - "max_input_tokens": 131072, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "realtime", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 4.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/realtime" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_vision": true, - "supports_web_search": true, - "gemini_audio_only_live": true, - "supports_reasoning": true - }, "gemini/gemini-3.1-flash-tts-preview": { "input_cost_per_token": 1e-06, "litellm_provider": "gemini", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ce531ba1fe1..5631220144e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -56012,73 +56012,6 @@ "rpm": 10, "gemini_audio_only_live": true }, - "gemini/gemini-3.8-live": { - "input_cost_per_audio_token": 3e-06, - "input_cost_per_image_token": 1e-06, - "input_cost_per_token": 7.5e-07, - "input_cost_per_video_per_second": 3.3333333333333335e-05, - "litellm_provider": "gemini", - "max_input_tokens": 131072, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "realtime", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 4.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/realtime" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_vision": true, - "supports_web_search": true, - "gemini_audio_only_live": true - }, - "gemini/gemini-3.8-live-extended-thinking": { - "input_cost_per_audio_token": 3e-06, - "input_cost_per_image_token": 1e-06, - "input_cost_per_token": 7.5e-07, - "input_cost_per_video_per_second": 3.3333333333333335e-05, - "litellm_provider": "gemini", - "max_input_tokens": 131072, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "realtime", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 4.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/realtime" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_vision": true, - "supports_web_search": true, - "gemini_audio_only_live": true, - "supports_reasoning": true - }, "gemini/gemini-3.1-flash-tts-preview": { "input_cost_per_token": 1e-06, "litellm_provider": "gemini", From 65a2e4f5039449c64c66507a8c63967c00004d0e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:38:52 +0000 Subject: [PATCH 034/187] test(cost): derive Nova cache expectations from registry and cite Gemini alias targets Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ..._cross_region_inference_profile_mapping.py | 22 +++++++++---------- .../llms/gemini/test_cost_calculator.py | 2 ++ 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index b567a8669b4..b12ca31c694 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -230,21 +230,21 @@ def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(prof @pytest.mark.parametrize( - "model,expected_cache_read", + "model", [ - ("amazon.nova-lite-v1:0", 1.5e-8), - ("us.amazon.nova-lite-v1:0", 1.5e-8), - ("amazon.nova-micro-v1:0", 8.75e-9), - ("us.amazon.nova-micro-v1:0", 8.75e-9), - ("amazon.nova-pro-v1:0", 2e-7), - ("us.amazon.nova-pro-v1:0", 2e-7), - ("us.amazon.nova-premier-v1:0", 6.25e-7), + "amazon.nova-lite-v1:0", + "us.amazon.nova-lite-v1:0", + "amazon.nova-micro-v1:0", + "us.amazon.nova-micro-v1:0", + "amazon.nova-pro-v1:0", + "us.amazon.nova-pro-v1:0", + "us.amazon.nova-premier-v1:0", ], ) -def test_bedrock_nova_cache_read_prices( - model, expected_cache_read, local_model_cost_map -): +def test_bedrock_nova_cache_read_prices(model, local_model_cost_map): model_info = litellm.model_cost[model] + expected_cache_read = model_info["cache_read_input_token_cost"] + assert 0 < expected_cache_read < model_info["input_cost_per_token"] usage = Usage( prompt_tokens=1_000, completion_tokens=100, diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index 8ad27411ce3..5eed11dff03 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -452,6 +452,8 @@ def test_map_traffic_type_to_service_tier( ) +# Alias targets are the `modelVersion` returned by +# POST https://generativelanguage.googleapis.com/v1beta/models/:generateContent on 2026-09-15 @pytest.mark.parametrize( "alias,target", [ From 5f0934309918b2b65c39c65f5eae3d1321e9ca3c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:50:01 +0000 Subject: [PATCH 035/187] test(bedrock): pin Nova cache-read prices with dated AWS Price List citation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ..._cross_region_inference_profile_mapping.py | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index b12ca31c694..b1e85c4f5ea 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -229,22 +229,25 @@ def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(prof assert "output_config" not in supported +# Cache-read prices are the `*-cache-read-input-tokens` usagetype rows of the AWS Price List API, us-east-1, +# https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json on 2026-09-15 @pytest.mark.parametrize( - "model", + "model,expected_cache_read", [ - "amazon.nova-lite-v1:0", - "us.amazon.nova-lite-v1:0", - "amazon.nova-micro-v1:0", - "us.amazon.nova-micro-v1:0", - "amazon.nova-pro-v1:0", - "us.amazon.nova-pro-v1:0", - "us.amazon.nova-premier-v1:0", + ("amazon.nova-lite-v1:0", 1.5e-8), + ("us.amazon.nova-lite-v1:0", 1.5e-8), + ("amazon.nova-micro-v1:0", 8.75e-9), + ("us.amazon.nova-micro-v1:0", 8.75e-9), + ("amazon.nova-pro-v1:0", 2e-7), + ("us.amazon.nova-pro-v1:0", 2e-7), + ("us.amazon.nova-premier-v1:0", 6.25e-7), ], ) -def test_bedrock_nova_cache_read_prices(model, local_model_cost_map): +def test_bedrock_nova_cache_read_prices( + model, expected_cache_read, local_model_cost_map +): model_info = litellm.model_cost[model] - expected_cache_read = model_info["cache_read_input_token_cost"] - assert 0 < expected_cache_read < model_info["input_cost_per_token"] + assert model_info["cache_read_input_token_cost"] == expected_cache_read usage = Usage( prompt_tokens=1_000, completion_tokens=100, From 6c8b9a7b0554482974fb3043875f550096e36de9 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 22:12:01 +0000 Subject: [PATCH 036/187] feat(keys): filter /key/list by active, expired, revoked or deleted status and serve deleted keys from /key/info Persist and expose the lifecycle of API keys so spend, audit and FinOps workflows can still resolve a key after it is revoked, expires or is deleted. /key/list?status= now accepts active, expired and revoked next to the existing deleted value. revoked means blocked=true, expired means not blocked with a past expiry, active is the rest, so the three values partition the live key table. deleted keeps reading the LiteLLM_DeletedVerificationToken archive. /key/info falls back to that archive when the key is no longer in the live table, running the same owner/team/org authorization check, and every response now carries a derived status field. The hashed token is still stripped. The Virtual Keys page gets a Status filter (URL-persisted) and a Deleted badge that shows when and by whom the key was deleted. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 108 +++++++-- .../test_key_management_endpoints.py | 218 ++++++++++++++++++ .../VirtualKeysPage/VirtualKeysTable.test.tsx | 54 +++++ .../VirtualKeysPage/VirtualKeysTable.tsx | 52 ++++- .../VirtualKeysPage/keyTableColumns.tsx | 7 + .../components/key_team_helpers/key_list.tsx | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 11 +- 7 files changed, 422 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ee8ae66ea11..7a363729541 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4166,7 +4166,10 @@ async def info_key_fn( Returns: - key: str - The key that was looked up, echoed back as it was passed in - - info: dict - The key's row, minus the hashed token + - info: dict - The key's row, minus the hashed token. Deleted keys are served from the + LiteLLM_DeletedVerificationToken archive and carry deleted_at / deleted_by + - status: "active" | "expired" | "revoked" | "deleted" - Derived from blocked, expires and + whether the row came from the archive - key_alias: str | None - User-friendly key alias - spend: float - Amount spent by the key. When budget_duration is set this covers only the current budget window, not the key's lifetime @@ -4220,10 +4223,15 @@ async def info_key_fn( hashed_key: str | None = key if key is not None: hashed_key = _hash_token_if_needed(token=key) - key_info = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( + live_key_info: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( where={"token": hashed_key}, include={"litellm_budget_table": True}, ) + key_info: Final = ( + live_key_info + if live_key_info is not None + else await _find_deleted_key_info(prisma_client=prisma_client, hashed_key=hashed_key) + ) if key_info is None: raise ProxyException( message="Key not found in database", @@ -4231,7 +4239,6 @@ async def info_key_fn( param="key", code=status.HTTP_404_NOT_FOUND, ) - if ( await _can_user_query_key_info( user_api_key_dict=user_api_key_dict, @@ -4245,38 +4252,46 @@ async def info_key_fn( detail=f"You are not allowed to access this key's info. Your role={user_api_key_dict.user_role}", ) ## REMOVE HASHED TOKEN INFO BEFORE RETURNING ## - try: - key_info = key_info.model_dump() - except Exception: - # if using pydantic v1 - key_info = key_info.dict() # pyright: ignore[reportDeprecated] # deliberate pydantic v1 fallback - key_token_hash: Final[str | None] = key_info.pop("token") + key_info_dict: Final = key_info.model_dump() + key_token_hash: Final[str | None] = key_info_dict.pop("token") + key_info_dict["status"] = ( + "deleted" if live_key_info is None else _derive_key_status(key_info_dict, now=datetime.now(timezone.utc)) + ) - model_max_budget = key_info.get("model_max_budget") or {} - budget_table: Final = key_info.get("litellm_budget_table") or {} + model_max_budget = key_info_dict.get("model_max_budget") or {} + budget_table: Final = key_info_dict.get("litellm_budget_table") or {} if not model_max_budget and isinstance(budget_table, dict): model_max_budget = budget_table.get("model_max_budget") or {} if model_max_budget and key_token_hash: - key_info["model_max_budget_usage"] = await _build_model_max_budget_usage( + key_info_dict["model_max_budget_usage"] = await _build_model_max_budget_usage( api_key_hash=key_token_hash, model_max_budget=model_max_budget, user_api_key_cache=model_max_budget_limiter.dual_cache, ) budget_limits_usage: Final = await _build_budget_limits_usage( - budget_limits=key_info.get("budget_limits"), + budget_limits=key_info_dict.get("budget_limits"), api_key_hash=key_token_hash, ) if budget_limits_usage is not None: - key_info["budget_limits_usage"] = budget_limits_usage + key_info_dict["budget_limits_usage"] = budget_limits_usage - # Attach object_permission if object_permission_id is set - key_info = await attach_object_permission_to_dict(key_info, prisma_client) - - return {"key": key, "info": key_info} + return {"key": key, "info": await attach_object_permission_to_dict(key_info_dict, prisma_client)} except Exception as e: raise handle_exception_on_proxy(e) +async def _find_deleted_key_info( + prisma_client: PrismaClient, hashed_key: str | None +) -> LiteLLM_DeletedVerificationToken | None: + archived_row: Final = await _deleted_verification_token_table(prisma_client).find_first( + where={"token": hashed_key}, + order={"deleted_at": "desc"}, + ) + if archived_row is None: + return None + return LiteLLM_DeletedVerificationToken.model_validate(archived_row.model_dump()) + + def _check_model_access_group(models: list[str] | None, llm_router: Router | None, premium_user: bool) -> Literal[True]: """ if is_model_access_group is True + is_wildcard_route is True, check if user is a premium user @@ -6216,6 +6231,25 @@ async def get_member_team_ids( VALID_EXPIRES_FILTER_VALUES: Final = frozenset({"active", "expired"}) +KeyStatus = Literal["active", "expired", "revoked", "deleted"] +VALID_STATUS_FILTER_VALUES: Final[frozenset[KeyStatus]] = frozenset({"active", "expired", "revoked", "deleted"}) + + +class _KeyStatusSource(BaseModel): + blocked: bool | None = None + expires: datetime | None = None + + +def _derive_key_status(row: Mapping[str, object], now: datetime) -> KeyStatus: + """Status of a live key row; mirrors the partition `_build_status_where_clause` applies at query time.""" + source: Final = _KeyStatusSource.model_validate(row) + if source.blocked is True: + return "revoked" + if source.expires is None: + return "active" + expires_utc: Final = source.expires if source.expires.tzinfo else source.expires.replace(tzinfo=timezone.utc) + return "expired" if expires_utc < now else "active" + @router.get( "/key/list", @@ -6252,7 +6286,10 @@ async def list_keys( ), sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), expand: list[str] | None = Query(None, description="Expand related objects (e.g. 'user')"), - status: str | None = Query(None, description="Filter by status (e.g. 'deleted')"), + status: str | None = Query( + None, + description="Filter by status: 'active' (not blocked, not expired), 'expired' (not blocked, past expiry), 'revoked' (blocked) or 'deleted' (archived keys). Omit to return live keys regardless of status.", + ), project_id: str | None = Query(None, description="Filter keys by project ID"), access_group_id: str | None = Query(None, description="Filter keys by access group ID"), agent_id: str | None = Query(None, description="Filter keys by agent ID"), @@ -6270,7 +6307,9 @@ async def list_keys( Parameters: expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information) - status: Optional[str] - Filter by status. Currently supports "deleted" to query deleted keys. + status: Optional[str] - Filter by status: "active", "expired", "revoked" (blocked) or "deleted". + "deleted" reads the LiteLLM_DeletedVerificationToken archive; the other values partition the + live key table, so every live key matches exactly one of them. Returns: { @@ -6292,11 +6331,10 @@ async def list_keys( verbose_proxy_logger.error("Database not connected") raise Exception("Database not connected") - # Validate status parameter - if status is not None and status != "deleted": + if status is not None and status not in VALID_STATUS_FILTER_VALUES: raise HTTPException( status_code=400, - detail={"error": "Invalid status value. Currently only 'deleted' is supported."}, + detail={"error": "Invalid status value. Supported: 'active', 'expired', 'revoked', 'deleted'."}, ) if isinstance(expires, str) and expires not in VALID_EXPIRES_FILTER_VALUES: @@ -6608,6 +6646,23 @@ def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str, return {"OR": [{"expires": None}, {"expires": {"gte": now}}]} +def _not_blocked_where_clause() -> dict[str, object]: + return {"OR": [{"blocked": None}, {"blocked": False}]} + + +def _build_status_where_clause(status_filter: str | None, now: datetime) -> dict[str, object] | None: + """Live-table clause for a status filter; None when the status needs no clause (deleted rows live elsewhere).""" + match status_filter: + case "revoked": + return {"blocked": True} + case "expired": + return {"AND": [_not_blocked_where_clause(), _build_expires_where_clause("expired", now)]} + case "active": + return {"AND": [_not_blocked_where_clause(), _build_expires_where_clause("active", now)]} + case _: + return None + + def _build_key_search_where(search: str) -> KeySearchWhere: search_where: Final[KeySearchWhere] = { "OR": ( @@ -6635,6 +6690,7 @@ def _build_key_filter_conditions( use_key_alias_substring_matching: bool = False, expires_filter: str | None = None, search: str | None = None, + status_filter: str | None = None, ) -> Mapping[str, object]: """Build filter conditions for key listing. @@ -6724,6 +6780,8 @@ def _build_key_filter_conditions( # Apply team_id, project_id and access_group_id as global AND filters so they # narrow results across all visibility conditions (own keys, team keys, etc.) + now: Final = datetime.now(timezone.utc) + status_where: Final = _build_status_where_clause(status_filter, now) global_filters: Final[tuple[Mapping[str, object], ...]] = ( *( ( @@ -6741,10 +6799,11 @@ def _build_key_filter_conditions( *(({"access_group_ids": {"hasSome": [access_group_id]}},) if access_group_id else ()), *(({"agent_id": agent_id},) if agent_id and isinstance(agent_id, str) else ()), *( - (_build_expires_where_clause(expires_filter, datetime.now(timezone.utc)),) + (_build_expires_where_clause(expires_filter, now),) if expires_filter is not None and expires_filter in VALID_EXPIRES_FILTER_VALUES else () ), + *((status_where,) if status_where is not None else ()), ) combined_where: Final[Mapping[str, object]] = {"AND": [where, *global_filters]} if global_filters else where verbose_proxy_logger.debug("Filter conditions: %s", combined_where) @@ -6817,6 +6876,7 @@ async def _list_key_helper( use_key_alias_substring_matching=use_key_alias_substring_matching, expires_filter=expires_filter, search=search, + status_filter=status, ) # Calculate skip for pagination diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 63055872aa1..852b8632846 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6017,6 +6017,224 @@ async def test_list_keys_with_invalid_status(): assert "deleted" in str(exc_info.value.message) +@pytest.mark.asyncio +@pytest.mark.parametrize("status_filter", ["active", "expired", "revoked"]) +async def test_list_keys_accepts_live_status_filters(monkeypatch, status_filter): + """LIT-1650: /key/list used to 400 on every status but "deleted"; the live statuses reach the helper.""" + from unittest.mock import Mock + + from litellm.proxy.management_endpoints import key_management_endpoints + + helper = AsyncMock(return_value={"keys": [], "total_count": 0, "current_page": 1, "total_pages": 0}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock()) + monkeypatch.setattr(key_management_endpoints, "_list_key_helper", helper) + await key_management_endpoints.list_keys( + request=Mock(), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + status=status_filter, + ) + + assert helper.await_args is not None + assert helper.await_args.kwargs["status"] == status_filter + + +def _status_filter_where(status_filter: str | None) -> Mapping[str, object]: + from litellm.proxy.management_endpoints.key_management_endpoints import _build_key_filter_conditions + + return _build_key_filter_conditions( + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=None, + status_filter=status_filter, + ) + + +def test_build_key_filter_conditions_status_filter_partitions_live_keys(): + """LIT-1650: active, expired and revoked are disjoint predicates over blocked + expires on the live table.""" + not_blocked = {"OR": [{"blocked": None}, {"blocked": False}]} + + revoked_where = _status_filter_where("revoked") + assert {"blocked": True} in revoked_where["AND"] + + expired_clause = next(clause for clause in _status_filter_where("expired")["AND"] if "AND" in clause) + assert expired_clause["AND"][0] == not_blocked + assert expired_clause["AND"][1]["AND"][0] == {"expires": {"not": None}} + assert "lt" in expired_clause["AND"][1]["AND"][1]["expires"] + + active_clause = next(clause for clause in _status_filter_where("active")["AND"] if "AND" in clause) + assert active_clause["AND"][0] == not_blocked + assert active_clause["AND"][1]["OR"][0] == {"expires": None} + assert "gte" in active_clause["AND"][1]["OR"][1]["expires"] + + +def test_build_key_filter_conditions_deleted_status_adds_no_live_clause(): + """Deleted rows live in the archive table, so the status must not narrow the live-table query.""" + assert _status_filter_where("deleted") == _status_filter_where(None) + + +@pytest.mark.asyncio +async def test_list_key_helper_revoked_status_filters_live_table_on_blocked(): + """LIT-1650: status="revoked" stays on the live table and narrows it to blocked keys.""" + mock_prisma_client = AsyncMock() + mock_find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + + await _list_key_helper( + prisma_client=mock_prisma_client, + page=1, + size=50, + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + return_full_object=True, + admin_team_ids=None, + include_created_by_keys=False, + status="revoked", + ) + + mock_prisma_client.db.litellm_deletedverificationtoken.find_many.assert_not_called() + where = mock_find_many.call_args.kwargs["where"] + assert {"blocked": True} in where["AND"] + + +def _archived_key_row(token: str, user_id: str) -> MagicMock: + row = MagicMock() + row.model_dump.return_value = { + "id": "archive-row-1", + "token": token, + "key_alias": "finops-2024", + "user_id": user_id, + "team_id": None, + "blocked": None, + "deleted_at": datetime(2024, 11, 15, 10, 0, tzinfo=timezone.utc), + "deleted_by": "admin-1", + } + return row + + +@pytest.mark.asyncio +async def test_info_key_fn_serves_deleted_key_from_archive(monkeypatch): + """LIT-1650: /key/info falls back to LiteLLM_DeletedVerificationToken and reports status="deleted".""" + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + hashed = "hashed_deleted_token" + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock( + return_value=_archived_key_row(hashed, "user-x") + ) + + result = await info_key_fn( + key=hashed, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"), + ) + + mock_prisma_client.db.litellm_deletedverificationtoken.find_first.assert_awaited_once() + assert mock_prisma_client.db.litellm_deletedverificationtoken.find_first.await_args.kwargs["where"] == { + "token": hashed + } + info = result["info"] + assert info["status"] == "deleted" + assert info["key_alias"] == "finops-2024" + assert info["deleted_by"] == "admin-1" + assert info["deleted_at"] is not None + assert "token" not in info + + +@pytest.mark.asyncio +async def test_info_key_fn_archived_key_keeps_owner_authorization(monkeypatch): + """An archived key is still scoped: a different internal user gets 403, the owner gets the record.""" + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + hashed = "hashed_deleted_token" + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock( + return_value=_archived_key_row(hashed, "owner-1") + ) + + with pytest.raises(ProxyException) as exc_info: + await info_key_fn( + key=hashed, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="someone-else", api_key="sk-other" + ), + ) + assert exc_info.value.code == "403" + + owner_result = await info_key_fn( + key=hashed, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="owner-1", api_key="sk-own"), + ) + assert owner_result["info"]["status"] == "deleted" + + +@pytest.mark.asyncio +async def test_info_key_fn_unknown_key_still_404s(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=None) + + with pytest.raises(ProxyException) as exc_info: + await info_key_fn( + key="hashed_missing", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"), + ) + assert exc_info.value.code == "404" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("blocked", "expires", "expected_status"), + [ + (True, None, "revoked"), + (True, "2020-01-01T00:00:00Z", "revoked"), + (False, "2020-01-01T00:00:00Z", "expired"), + (None, datetime(2020, 1, 1, tzinfo=timezone.utc), "expired"), + (False, None, "active"), + (None, "2999-01-01T00:00:00Z", "active"), + ], +) +async def test_info_key_fn_reports_live_key_status(monkeypatch, blocked, expires, expected_status): + """LIT-1650: live keys carry the same status vocabulary /key/list filters on.""" + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + live_row = MagicMock(spec=LiteLLM_VerificationToken) + live_row.model_dump.return_value = { + "token": "hashed_live", + "user_id": "user-x", + "team_id": None, + "object_permission_id": None, + "blocked": blocked, + "expires": expires, + } + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=live_row) + + result = await info_key_fn( + key="hashed_live", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"), + ) + + assert result["info"]["status"] == expected_status + mock_prisma_client.db.litellm_deletedverificationtoken.find_first.assert_not_called() + + @pytest.mark.asyncio async def test_list_keys_non_admin_user_id_auto_set(): """ diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 617b9209a41..fa09b2c1b5b 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -638,6 +638,23 @@ describe("server-side filtering – the LIT-4080 regression guard", () => { }); }); + it("threads the Status drawer filter into the useKeys query and the URL", async () => { + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + + openFilters(); + const user = userEvent.setup(); + await chooseSelectOption(user, await screen.findByRole("combobox", { name: "Status" }), "Revoked (blocked)"); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ status: "revoked" })); + }); + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_status")).toBe("revoked"); + }); + }); + it("sends the search box as the combined alias-or-ID search rather than the key-alias filter", async () => { renderWithProviders(); @@ -745,6 +762,25 @@ describe("Status column reflects blocked / expiry / scim metadata", () => { expect(screen.queryByText(/Blocked by SCIM/i)).not.toBeInTheDocument(); }); + it("renders Deleted for an archived key, even when the archived row was also blocked", async () => { + mockUseKeys.mockReturnValue( + keysResult([ + { ...mockKey, blocked: true, metadata: {}, deleted_at: "2024-11-15T10:00:00Z", deleted_by: "admin-1" }, + ]), + ); + + renderWithProviders(); + + const tag = await screen.findByTestId(`key-status-${mockKey.token_id}`); + expect(tag).toHaveTextContent("Deleted"); + + const user = userEvent.setup(); + await user.hover(tag); + await waitFor(() => { + expect(screen.getByText(/by admin-1/)).toBeInTheDocument(); + }); + }); + it("marks a SCIM-blocked key with the SCIM tooltip reason", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, blocked: true, metadata: { scim_blocked: true } }])); @@ -790,6 +826,24 @@ describe("table state lives in the URL so it survives leaving and returning to t expect(screen.getByTestId("filter-chip-team_id")).toHaveTextContent("Test Team"); }); + it("restores the status filter from the URL and sends it to /key/list", async () => { + renderWithProviders(, { searchParams: { filter_status: "deleted" } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ status: "deleted" })); + }); + expect(screen.getByTestId("filter-chip-status")).toHaveTextContent("Deleted"); + }); + + it("ignores a hand-edited status the backend would reject instead of 400ing the page", async () => { + renderWithProviders(, { searchParams: { filter_status: "bogus" } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ status: undefined })); + }); + expect(screen.queryByTestId("filter-chip-status")).not.toBeInTheDocument(); + }); + it("writes the search term to the URL", async () => { const onUrlUpdate = vi.fn(); renderWithProviders(, { onUrlUpdate }); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 907ee28bd05..1f52bdd7335 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -14,6 +14,7 @@ import { import { SearchSelect } from "@/components/shared/SearchSelect"; import { PageHeader } from "@/components/shared/PageHeader"; import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; import { KeyRound } from "lucide-react"; @@ -28,7 +29,7 @@ interface VirtualKeysTableProps { headerActions?: React.ReactNode; } -const FILTER_COLUMNS = ["team_id", "org_id", "user_id", "key_hash"] as const; +const FILTER_COLUMNS = ["team_id", "org_id", "user_id", "key_hash", "status"] as const; type FilterColumn = (typeof FILTER_COLUMNS)[number]; const FILTER_LABELS: Record = { @@ -36,8 +37,28 @@ const FILTER_LABELS: Record = { org_id: "Organization", user_id: "User ID", key_hash: "Key ID", + status: "Status", }; +const KEY_STATUS_VALUES = ["active", "expired", "revoked", "deleted"] as const; +type KeyStatusFilter = (typeof KEY_STATUS_VALUES)[number]; +const ALL_STATUSES = "all"; + +const KEY_STATUS_LABELS: Record = { + active: "Active", + expired: "Expired", + revoked: "Revoked (blocked)", + deleted: "Deleted", +}; + +const STATUS_FILTER_ITEMS = [ + { value: ALL_STATUSES, label: "All statuses" }, + ...KEY_STATUS_VALUES.map((value) => ({ value, label: KEY_STATUS_LABELS[value] })), +]; + +const isKeyStatusFilter = (value: string): value is KeyStatusFilter => + (KEY_STATUS_VALUES as readonly string[]).includes(value); + const DEFAULT_SORT_BY = "created_at"; const DEFAULT_SORT_ORDER = "desc"; const DEFAULT_PAGE_SIZE = 50; @@ -65,6 +86,7 @@ const TABLE_STATE = { filter_org: parseAsString.withDefault(""), filter_user: parseAsString.withDefault(""), filter_key_id: parseAsString.withDefault(""), + filter_status: parseAsString.withDefault(""), }; const toSortOrder = (active: SortingState[number]): "asc" | "desc" => (active.desc ? "desc" : "asc"); @@ -96,15 +118,16 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { () => ({ pageIndex: tableState.page - 1, pageSize: tableState.page_size }), [tableState.page, tableState.page_size], ); - const { filter_team, filter_org, filter_user, filter_key_id } = tableState; + const { filter_team, filter_org, filter_user, filter_key_id, filter_status } = tableState; const appliedFilters = useMemo( () => ({ team_id: filter_team.trim(), org_id: filter_org.trim(), user_id: filter_user.trim(), key_hash: filter_key_id.trim(), + status: isKeyStatusFilter(filter_status) ? filter_status : "", }), - [filter_team, filter_org, filter_user, filter_key_id], + [filter_team, filter_org, filter_user, filter_key_id, filter_status], ); const columnFilters = useMemo( () => @@ -121,6 +144,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { search: searchQuery.trim() || undefined, userID: appliedFilters.user_id || undefined, keyHash: appliedFilters.key_hash || undefined, + status: appliedFilters.status || undefined, sortBy, sortOrder: tableState.sort_order, expand: "user", @@ -164,6 +188,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { filter_org: filterValue(next, "org_id"), filter_user: filterValue(next, "user_id"), filter_key_id: filterValue(next, "key_hash"), + filter_status: filterValue(next, "status"), page: null, }; void setTableState(nextFilters); @@ -233,6 +258,9 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { if (columnId === "org_id") { return organizations.find((org) => org.organization_id === raw)?.organization_alias || raw; } + if (columnId === "status" && isKeyStatusFilter(raw)) { + return KEY_STATUS_LABELS[raw]; + } return raw; }, [allTeams, organizations], @@ -340,6 +368,24 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { placeholder="Enter Key ID…" /> + + + )} diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index 6eea77ae827..ff608365500 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -43,6 +43,13 @@ export const KEY_TABLE_SORT_FIELDS: readonly string[] = [ ]; const getKeyStatus = (key: KeyResponse): KeyStatus => { + if (key.deleted_at) { + return { + tone: "neutral", + label: "Deleted", + tooltip: `Deleted ${new Date(key.deleted_at).toLocaleString()}${key.deleted_by ? ` by ${key.deleted_by}` : ""}. Kept for audit and spend history; requests using this key are rejected.`, + }; + } if (key.blocked === true) { const isScimBlocked = (key.metadata as Record | null | undefined)?.scim_blocked === true; return { diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index eadbca87140..b42bef04be8 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -64,6 +64,8 @@ export interface KeyResponse { model_max_budget_usage?: Record | null; soft_budget_cooldown: boolean; blocked: boolean; + deleted_at?: string | null; + deleted_by?: string | null; litellm_budget_table: Record; organization_id: string | null; org_id?: string | null; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 4c16a8613eb..542d491e2d7 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7859,7 +7859,10 @@ export interface paths { * * Returns: * - key: str - The key that was looked up, echoed back as it was passed in - * - info: dict - The key's row, minus the hashed token + * - info: dict - The key's row, minus the hashed token. Deleted keys are served from the + * LiteLLM_DeletedVerificationToken archive and carry deleted_at / deleted_by + * - status: "active" | "expired" | "revoked" | "deleted" - Derived from blocked, expires and + * whether the row came from the archive * - key_alias: str | None - User-friendly key alias * - spend: float - Amount spent by the key. When budget_duration is set this covers only the * current budget window, not the key's lifetime @@ -7917,7 +7920,9 @@ export interface paths { * * Parameters: * expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information) - * status: Optional[str] - Filter by status. Currently supports "deleted" to query deleted keys. + * status: Optional[str] - Filter by status: "active", "expired", "revoked" (blocked) or "deleted". + * "deleted" reads the LiteLLM_DeletedVerificationToken archive; the other values partition the + * live key table, so every live key matches exactly one of them. * * Returns: * { @@ -51185,7 +51190,7 @@ export interface operations { sort_order?: string; /** @description Expand related objects (e.g. 'user') */ expand?: string[] | null; - /** @description Filter by status (e.g. 'deleted') */ + /** @description Filter by status: 'active' (not blocked, not expired), 'expired' (not blocked, past expiry), 'revoked' (blocked) or 'deleted' (archived keys). Omit to return live keys regardless of status. */ status?: string | null; /** @description Filter keys by project ID */ project_id?: string | null; From f6f782ff67168b6e52e8e03828c5c6270c850944 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 22:56:37 +0000 Subject: [PATCH 037/187] feat(s3): add s3_log_prompts_only option to log prompts without responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/integrations/callback_configs.json | 6 + litellm/integrations/custom_logger.py | 1 + litellm/integrations/s3.py | 69 ++++++--- litellm/integrations/s3_v2.py | 16 ++- litellm/proxy/_types.py | 1 + tests/test_litellm/integrations/test_s3.py | 136 +++++++++++++++++- tests/test_litellm/integrations/test_s3_v2.py | 130 +++++++++++++++++ .../src/components/settings.test.tsx | 106 ++++++++++++++ .../src/components/settings.tsx | 47 +++++- 10 files changed, 485 insertions(+), 28 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 1dbb8a842fb..3268c871ca6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -53,6 +53,7 @@ S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64 S3_PREFIX_DIGEST_CHARS: Final = 16 # s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024 +S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY" MAX_FILE_LIST_LIMIT: Final = 10000 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 85bfcc6e7ed..6806188c97c 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -446,6 +446,12 @@ "ui_name": "S3 Path Prefix", "description": "Path prefix within the bucket for organizing logs", "required": false + }, + "s3_log_prompts_only": { + "type": "boolean", + "ui_name": "Log Prompts Only", + "description": "Log request messages to S3 but drop the model response from each logged object", + "required": false } }, "description": "S3 Bucket (AWS) Logging Integration" diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 70d2f3ae5c3..5b5261fab6b 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -118,6 +118,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac alias_map: Final = { "langfuse_otel": "langfuse", + "s3_v2": "s3", } lookup_name: Final = alias_map.get(normalized_name, normalized_name) diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index 8ce461eea5b..796784fb993 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -2,19 +2,42 @@ # On success + failure, log events to Supabase import hashlib +import os +from collections.abc import Mapping from datetime import datetime from typing import Final, cast +from pydantic import TypeAdapter, ValidationError + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES, MAX_S3_OBJECT_KEY_BYTES, S3_BOUNDED_OBJECT_KEY_HEAD_BYTES, + S3_LOG_PROMPTS_ONLY_ENV_VAR, S3_PREFIX_DIGEST_CHARS, ) from litellm.types.utils import StandardLoggingPayload +_S3_LOG_PROMPTS_ONLY: Final = TypeAdapter(bool) + + +def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] | None = None) -> bool: + env: Final = os.environ if environ is None else environ + raw: Final = env.get(S3_LOG_PROMPTS_ONLY_ENV_VAR) if configured is None else configured + if raw is None or raw == "": + return False + try: + return _S3_LOG_PROMPTS_ONLY.validate_python(raw.strip() if isinstance(raw, str) else raw) + except ValidationError: + verbose_logger.warning("s3 logging: s3_log_prompts_only=%r is not a boolean, logging prompts only", raw) + return True + + +def prompts_only_payload(payload: StandardLoggingPayload) -> StandardLoggingPayload: + return {**payload, "response": None} + class S3Logger: # Class variables or attributes @@ -33,6 +56,7 @@ class S3Logger: s3_config=None, s3_server_side_encryption: str | None = None, s3_sse_kms_key_id: str | None = None, + s3_log_prompts_only: bool | None = None, **kwargs, ): import boto3 @@ -41,29 +65,30 @@ class S3Logger: verbose_logger.debug("in init s3 logger - s3_callback_params %s", litellm.s3_callback_params) s3_use_team_prefix = False + params: Final = { + key: litellm.get_secret(value) if isinstance(value, str) and value.startswith("os.environ/") else value + for key, value in (litellm.s3_callback_params or {}).items() + } if litellm.s3_callback_params is not None: - # read in .env variables - example os.environ/AWS_BUCKET_NAME - for key, value in litellm.s3_callback_params.items(): - if isinstance(value, str) and value.startswith("os.environ/"): - litellm.s3_callback_params[key] = litellm.get_secret(value) - # now set s3 params from litellm.s3_logger_params - s3_bucket_name = litellm.s3_callback_params.get("s3_bucket_name") - s3_region_name = litellm.s3_callback_params.get("s3_region_name") - s3_api_version = litellm.s3_callback_params.get("s3_api_version") - s3_use_ssl = litellm.s3_callback_params.get("s3_use_ssl", True) - s3_verify = litellm.s3_callback_params.get("s3_verify") - s3_endpoint_url = litellm.s3_callback_params.get("s3_endpoint_url") - s3_aws_access_key_id = litellm.s3_callback_params.get("s3_aws_access_key_id") - s3_aws_secret_access_key = litellm.s3_callback_params.get("s3_aws_secret_access_key") - s3_aws_session_token = litellm.s3_callback_params.get("s3_aws_session_token") - s3_config = litellm.s3_callback_params.get("s3_config") - s3_path = litellm.s3_callback_params.get("s3_path") - s3_server_side_encryption = litellm.s3_callback_params.get("s3_server_side_encryption") - s3_sse_kms_key_id = litellm.s3_callback_params.get("s3_sse_kms_key_id") - # done reading litellm.s3_callback_params - s3_use_team_prefix = bool(litellm.s3_callback_params.get("s3_use_team_prefix", False)) + s3_bucket_name = params.get("s3_bucket_name") + s3_region_name = params.get("s3_region_name") + s3_api_version = params.get("s3_api_version") + s3_use_ssl = params.get("s3_use_ssl", True) + s3_verify = params.get("s3_verify") + s3_endpoint_url = params.get("s3_endpoint_url") + s3_aws_access_key_id = params.get("s3_aws_access_key_id") + s3_aws_secret_access_key = params.get("s3_aws_secret_access_key") + s3_aws_session_token = params.get("s3_aws_session_token") + s3_config = params.get("s3_config") + s3_path = params.get("s3_path") + s3_server_side_encryption = params.get("s3_server_side_encryption") + s3_sse_kms_key_id = params.get("s3_sse_kms_key_id") + s3_use_team_prefix = bool(params.get("s3_use_team_prefix", False)) self.s3_use_team_prefix = s3_use_team_prefix + self.s3_log_prompts_only: object = ( + params.get("s3_log_prompts_only") if s3_log_prompts_only is None else s3_log_prompts_only + ) self.bucket_name = s3_bucket_name self.s3_path = s3_path self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params( @@ -144,7 +169,9 @@ class S3Logger: from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - payload_str: Final = safe_dumps(payload) + payload_str: Final = safe_dumps( + prompts_only_payload(payload) if resolve_s3_log_prompts_only(self.s3_log_prompts_only) else payload + ) print_verbose(f"\ns3 Logger - Logging payload = {payload_str}") diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 972ac79e306..826f55cc798 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -21,6 +21,8 @@ from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_S from litellm.integrations.s3 import ( get_s3_object_download_filename, get_s3_object_key, + prompts_only_payload, + resolve_s3_log_prompts_only, resolve_sse_params, ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix @@ -68,6 +70,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_virtual_hosted_style: bool = False, s3_server_side_encryption: str | None = None, s3_sse_kms_key_id: str | None = None, + s3_log_prompts_only: bool | None = None, s3_callback_params_override: dict | None = None, **kwargs, ): @@ -108,6 +111,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_virtual_hosted_style=s3_use_virtual_hosted_style, s3_server_side_encryption=s3_server_side_encryption, s3_sse_kms_key_id=s3_sse_kms_key_id, + s3_log_prompts_only=s3_log_prompts_only, ) verbose_logger.debug("s3 logger using endpoint url %s", s3_endpoint_url) @@ -163,6 +167,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_virtual_hosted_style: bool = False, s3_server_side_encryption: str | None = None, s3_sse_kms_key_id: str | None = None, + s3_log_prompts_only: bool | None = None, params_source: dict | None = None, ): """ @@ -212,6 +217,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): bool(params.get("s3_use_virtual_hosted_style", False)) or s3_use_virtual_hosted_style ) + self.s3_log_prompts_only: object = ( + params.get("s3_log_prompts_only") if s3_log_prompts_only is None else s3_log_prompts_only + ) + self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params( params.get("s3_server_side_encryption") or s3_server_side_encryption, params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id, @@ -489,8 +498,13 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_object_download_filename: Final = get_s3_object_download_filename(start_time, standard_logging_payload["id"]) + payload: Final = ( + prompts_only_payload(standard_logging_payload) + if resolve_s3_log_prompts_only(self.s3_log_prompts_only) + else standard_logging_payload + ) return s3BatchLoggingElement( - payload=dict(standard_logging_payload), + payload=dict(payload), s3_object_key=s3_object_key, s3_object_download_filename=s3_object_download_filename, ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ad55fa5d2be..e426e83bbe7 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3708,6 +3708,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase): "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME", + "S3_LOG_PROMPTS_ONLY", ], ) diff --git a/tests/test_litellm/integrations/test_s3.py b/tests/test_litellm/integrations/test_s3.py index 58b15b79e76..ba8d575c1b7 100644 --- a/tests/test_litellm/integrations/test_s3.py +++ b/tests/test_litellm/integrations/test_s3.py @@ -1,16 +1,24 @@ +import copy +import json from datetime import datetime from unittest.mock import MagicMock, patch +import pytest + import litellm from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES, MAX_S3_OBJECT_KEY_BYTES -from litellm.integrations.s3 import S3Logger +from litellm.integrations.s3 import S3Logger, prompts_only_payload, resolve_s3_log_prompts_only TEST_KMS_KEY_ARN = "arn:aws:kms:us-east-1:111122223333:key/test-key-id" +TEST_MESSAGES = [{"role": "user", "content": "Reply with exactly the word PINEAPPLE."}] +TEST_RESPONSE = {"choices": [{"message": {"role": "assistant", "content": "PINEAPPLE"}}]} def _standard_logging_payload(response_id: str = "chatcmpl-test-id") -> dict: return { "id": response_id, + "messages": copy.deepcopy(TEST_MESSAGES), + "response": copy.deepcopy(TEST_RESPONSE), "metadata": {"user_api_key_team_alias": None}, } @@ -22,7 +30,9 @@ def _log_event_kwargs(response_id: str = "chatcmpl-test-id") -> dict: } -def _run_log_event(callback_params: dict, response_id: str = "chatcmpl-test-id") -> MagicMock: +def _run_log_event( + callback_params: dict, response_id: str = "chatcmpl-test-id", log_kwargs: dict | None = None +) -> MagicMock: original = litellm.s3_callback_params litellm.s3_callback_params = callback_params try: @@ -31,7 +41,7 @@ def _run_log_event(callback_params: dict, response_id: str = "chatcmpl-test-id") mock_boto3_client.return_value = mock_s3_client logger = S3Logger() logger.log_event( - kwargs=_log_event_kwargs(response_id), + kwargs=_log_event_kwargs(response_id) if log_kwargs is None else log_kwargs, response_obj={"id": response_id}, start_time=datetime(2026, 7, 30, 12, 0, 0), end_time=datetime(2026, 7, 30, 12, 0, 1), @@ -182,3 +192,123 @@ def test_put_object_keeps_the_configured_path_intact_when_only_the_id_has_to_shr key = mock_s3_client.put_object.call_args.kwargs["Key"] assert key.startswith(long_path + "/2026-07-30/") assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES + + +def _uploaded_body(mock_s3_client: MagicMock) -> dict: + return json.loads(mock_s3_client.put_object.call_args.kwargs["Body"]) + + +def test_log_event_prompts_only_drops_response_and_keeps_messages(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) + log_kwargs = _log_event_kwargs() + original_payload = copy.deepcopy(log_kwargs["standard_logging_object"]) + + mock_s3_client = _run_log_event( + {"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1", "s3_log_prompts_only": True}, + log_kwargs=log_kwargs, + ) + + body = _uploaded_body(mock_s3_client) + assert body["messages"] == TEST_MESSAGES + assert body["response"] is None + assert body["id"] == "chatcmpl-test-id" + assert log_kwargs["standard_logging_object"] == original_payload + + +def test_log_event_default_keeps_response(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) + + mock_s3_client = _run_log_event({"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1"}) + + body = _uploaded_body(mock_s3_client) + assert body["response"] == TEST_RESPONSE + assert body["messages"] == TEST_MESSAGES + + +def test_log_event_reads_prompts_only_env_var_at_log_time(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) + original = litellm.s3_callback_params + litellm.s3_callback_params = {"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1"} + try: + with patch("boto3.client") as mock_boto3_client: + mock_s3_client = MagicMock() + mock_boto3_client.return_value = mock_s3_client + logger = S3Logger() + monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true") + logger.log_event( + kwargs=_log_event_kwargs(), + response_obj={"id": "chatcmpl-test-id"}, + start_time=datetime(2026, 7, 30, 12, 0, 0), + end_time=datetime(2026, 7, 30, 12, 0, 1), + print_verbose=lambda *args, **kwargs: None, + ) + finally: + litellm.s3_callback_params = original + + body = _uploaded_body(mock_s3_client) + assert body["response"] is None + assert body["messages"] == TEST_MESSAGES + + +def test_log_event_explicit_false_param_beats_env_var(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true") + + mock_s3_client = _run_log_event( + {"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1", "s3_log_prompts_only": False} + ) + + assert _uploaded_body(mock_s3_client)["response"] == TEST_RESPONSE + + +def test_s3_logger_init_does_not_mutate_global_callback_params(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MY_S3_BUCKET", "resolved-bucket") + callback_params = {"s3_bucket_name": "os.environ/MY_S3_BUCKET", "s3_region_name": "us-east-1"} + snapshot = copy.deepcopy(callback_params) + original = litellm.s3_callback_params + litellm.s3_callback_params = callback_params + try: + with patch("boto3.client"): + logger = S3Logger() + finally: + litellm.s3_callback_params = original + + assert logger.bucket_name == "resolved-bucket" + assert callback_params == snapshot + + +@pytest.mark.parametrize( + "configured,env_value,expected", + [ + (True, None, True), + (False, "true", False), + ("true", None, True), + ("False", "true", False), + ("1", None, True), + ("0", None, False), + (" yes ", None, True), + (None, None, False), + (None, "true", True), + (None, "false", False), + (None, "", False), + ("", "true", False), + ], +) +def test_resolve_s3_log_prompts_only(configured: object, env_value: str | None, expected: bool): + environ = {} if env_value is None else {"S3_LOG_PROMPTS_ONLY": env_value} + assert resolve_s3_log_prompts_only(configured, environ) is expected + + +def test_resolve_s3_log_prompts_only_unparseable_value_fails_toward_prompts_only(): + assert resolve_s3_log_prompts_only("enabled", {}) is True + + +def test_prompts_only_payload_returns_copy_with_response_cleared(): + payload = _standard_logging_payload() + snapshot = copy.deepcopy(payload) + + stripped = prompts_only_payload(payload) + + assert stripped["response"] is None + assert stripped["messages"] == TEST_MESSAGES + assert stripped is not payload + assert payload == snapshot diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 08d37297ab1..1179aa7e409 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1,4 +1,6 @@ import asyncio +import copy +import json import re import sys import textwrap @@ -10,6 +12,7 @@ from unittest.mock import AsyncMock, MagicMock, call, patch import httpx import pytest +import respx from litellm.integrations.s3_v2 import S3Logger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -2310,3 +2313,130 @@ def _s3_logger_for_region(region_name: str) -> S3Logger: ) def test_build_object_url_uses_partition_dns_suffix(region_name: str, expected_url: str) -> None: assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url + + +def _prompts_only_logger(**kwargs) -> S3Logger: + return S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + **kwargs, + ) + + +def _chat_payload() -> dict: + return { + "id": "chatcmpl-prompts-only", + "messages": [{"role": "user", "content": "Reply with exactly the word PINEAPPLE."}], + "response": {"choices": [{"message": {"role": "assistant", "content": "PINEAPPLE"}}]}, + "metadata": {"user_api_key_team_alias": None}, + } + + +async def _queued_body_via_async_upload(logger: S3Logger, log_event) -> dict: + payload = _chat_payload() + original = copy.deepcopy(payload) + await log_event( + kwargs={"standard_logging_object": payload}, + response_obj=None, + start_time=datetime(2026, 7, 30, 12, 0, 0), + end_time=datetime(2026, 7, 30, 12, 0, 1), + ) + assert payload == original, "the caller's standard_logging_object must not be mutated" + (element,) = logger.log_queue + + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put.return_value = response + await logger.async_upload_data_to_s3(element) + return json.loads(logger.async_httpx_client.put.call_args.kwargs["data"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_name", ["async_log_success_event", "async_log_failure_event"]) +async def test_prompts_only_drops_response_but_keeps_messages_in_uploaded_object(monkeypatch, event_name): + import litellm + + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": True}) + logger = _prompts_only_logger() + + body = await _queued_body_via_async_upload(logger, getattr(logger, event_name)) + + assert body["messages"] == _chat_payload()["messages"] + assert body["response"] is None + assert body["id"] == "chatcmpl-prompts-only" + + +@pytest.mark.asyncio +async def test_prompts_only_default_off_keeps_response_in_uploaded_object(monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "s3_callback_params", {}) + monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) + logger = _prompts_only_logger() + + body = await _queued_body_via_async_upload(logger, logger.async_log_success_event) + + assert body["response"] == _chat_payload()["response"] + assert body["messages"] == _chat_payload()["messages"] + + +@pytest.mark.asyncio +async def test_prompts_only_explicit_false_in_params_beats_env_var(monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": False}) + monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true") + logger = _prompts_only_logger() + + body = await _queued_body_via_async_upload(logger, logger.async_log_success_event) + + assert body["response"] == _chat_payload()["response"] + + +@pytest.mark.asyncio +async def test_prompts_only_env_var_applies_when_param_unset(monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "s3_callback_params", {}) + logger = _prompts_only_logger() + monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true") + + body = await _queued_body_via_async_upload(logger, logger.async_log_success_event) + + assert body["response"] is None + assert body["messages"] == _chat_payload()["messages"] + + +@respx.mock +def test_prompts_only_constructor_kwarg_applies_to_sync_upload(monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "s3_callback_params", {}) + monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) + logger = _prompts_only_logger(s3_log_prompts_only=True) + payload = _chat_payload() + + element = logger.create_s3_batch_logging_element( + start_time=datetime(2026, 7, 30, 12, 0, 0), + standard_logging_payload=payload, + ) + assert element is not None + assert payload["response"] == _chat_payload()["response"] + + put_route = respx.put(url__regex=r"https://test-bucket\.s3\..*").mock(return_value=httpx.Response(200)) + logger.upload_data_to_s3(element) + + body = json.loads(put_route.calls.last.request.content) + assert body["response"] is None + assert body["messages"] == _chat_payload()["messages"] + + +@pytest.mark.parametrize("callback_name", ["s3", "s3_v2"]) +def test_prompts_only_toggle_is_exposed_to_admin_ui_for_both_s3_callbacks(callback_name): + from litellm.integrations.custom_logger import CustomLogger + + assert "S3_LOG_PROMPTS_ONLY" in CustomLogger.get_callback_env_vars(callback_name) diff --git a/ui/litellm-dashboard/src/components/settings.test.tsx b/ui/litellm-dashboard/src/components/settings.test.tsx index 08cb9550646..c24fa00438a 100644 --- a/ui/litellm-dashboard/src/components/settings.test.tsx +++ b/ui/litellm-dashboard/src/components/settings.test.tsx @@ -302,6 +302,112 @@ describe("Settings", () => { }); }); + const mockS3Callback = (variables: Record, callbackName = "s3") => { + mockGetCallbacksCall.mockResolvedValue({ + callbacks: [{ name: callbackName, variables }], + available_callbacks: { + s3: { + litellm_callback_name: "s3", + litellm_callback_params: [ + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_REGION_NAME", + "S3_LOG_PROMPTS_ONLY", + ], + ui_callback_name: "s3 Bucket (AWS)", + }, + }, + alerts: [], + }); + mockGetCallbackConfigsCall.mockResolvedValue([ + { + id: "s3", + displayName: "S3", + dynamic_params: { + s3_bucket_name: { type: "text", ui_name: "S3 Bucket Name", required: false }, + s3_log_prompts_only: { type: "boolean", ui_name: "Log Prompts Only", required: false }, + }, + }, + ]); + }; + + const openS3EditModal = async (callbackName = "s3") => { + const user = userEvent.setup(); + render(); + await user.click(await screen.findByTestId(`callback-actions-${callbackName}-success`)); + await user.click(await screen.findByTestId("callback-action-edit")); + return user; + }; + + it("should render a saved boolean dynamic param as a checked switch and post false when toggled off", async () => { + mockS3Callback({ S3_LOG_PROMPTS_ONLY: "true" }); + const user = await openS3EditModal(); + + const promptsOnlySwitch = await screen.findByRole("switch", { name: "Log Prompts Only" }); + expect(promptsOnlySwitch).toBeChecked(); + + await user.click(promptsOnlySwitch); + expect(promptsOnlySwitch).not.toBeChecked(); + await user.click(within(screen.getByRole("dialog")).getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(vi.mocked(setCallbacksCall)).toHaveBeenCalledWith( + "token", + expect.objectContaining({ + environment_variables: expect.objectContaining({ callback: "s3", s3_log_prompts_only: "false" }), + }), + ); + }); + }); + + it("should render an unset boolean dynamic param as an unchecked switch and post true when toggled on", async () => { + mockS3Callback({ S3_LOG_PROMPTS_ONLY: null }); + const user = await openS3EditModal(); + + const promptsOnlySwitch = await screen.findByRole("switch", { name: "Log Prompts Only" }); + expect(promptsOnlySwitch).not.toBeChecked(); + + await user.click(promptsOnlySwitch); + await user.click(within(screen.getByRole("dialog")).getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(vi.mocked(setCallbacksCall)).toHaveBeenCalledWith( + "token", + expect.objectContaining({ + environment_variables: expect.objectContaining({ callback: "s3", s3_log_prompts_only: "true" }), + }), + ); + }); + }); + + it.each(["True", "1"])("should render a boolean dynamic param stored as %s as a checked switch", async (stored) => { + mockS3Callback({ S3_LOG_PROMPTS_ONLY: stored }); + await openS3EditModal(); + + expect(await screen.findByRole("switch", { name: "Log Prompts Only" })).toBeChecked(); + }); + + it("should resolve the s3_v2 callback to the s3 dynamic params and post under the s3_v2 name", async () => { + mockS3Callback({ S3_LOG_PROMPTS_ONLY: null }, "s3_v2"); + const user = await openS3EditModal("s3_v2"); + + const promptsOnlySwitch = await screen.findByRole("switch", { name: "Log Prompts Only" }); + expect(promptsOnlySwitch).not.toBeChecked(); + + await user.click(promptsOnlySwitch); + await user.click(within(screen.getByRole("dialog")).getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(vi.mocked(setCallbacksCall)).toHaveBeenCalledWith( + "token", + expect.objectContaining({ + environment_variables: expect.objectContaining({ callback: "s3_v2", s3_log_prompts_only: "true" }), + litellm_settings: { success_callback: ["s3_v2"] }, + }), + ); + }); + }); + it("should send the typed webhook url for an alert type when the alerting tab is saved", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index e549770af6e..c0daeac3b72 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -67,19 +67,20 @@ const DynamicParamsFields: React.FC = ({ params, callb return null; } + const callbackConfig = findCallbackConfig(callbackConfigs, selectedCallback); return (
{params.map((param) => { - const callbackConfig = callbackConfigs.find((config) => config.id === selectedCallback); const paramConfig = callbackConfig?.dynamic_params?.[param] || {}; const paramType = paramConfig.type || "text"; const fieldLabel = paramConfig.ui_name || param.replace(/_/g, " ").replace(/\b\w/g, (l) => l.toUpperCase()); const isRequired = paramConfig.required || false; const selectOptions: string[] = Array.isArray(paramConfig.options) ? paramConfig.options : []; const isSelect = paramType === "select" && selectOptions.length > 0; + const isBoolean = paramType === "boolean"; const fieldId = `${fieldIdPrefix}-${param}`; const validationRules = isRequired ? { required: `Please enter the ${fieldLabel.toLowerCase()}` } : undefined; - const registration = isSelect ? undefined : register(param, validationRules); + const registration = isSelect || isBoolean ? undefined : register(param, validationRules); return ( @@ -111,7 +112,22 @@ const DynamicParamsFields: React.FC = ({ params, callb )} /> )} + {isBoolean && ( + ( + field.onChange(checked ? "true" : "false")} + onBlur={field.onBlur} + /> + )} + /> + )} {!isSelect && + !isBoolean && (paramType === "password" ? ( = ({ ); }; +const CALLBACK_CONFIG_ALIASES: Record = { s3_v2: "s3" }; + +interface DynamicParamConfig { + type?: string; + ui_name?: string; + required?: boolean; + options?: string[]; +} + +interface CallbackConfigWithParams { + id: string; + dynamic_params?: Record; +} + +const findCallbackConfig = ( + callbackConfigs: readonly CallbackConfigWithParams[], + callbackName: string | null, +): CallbackConfigWithParams | undefined => { + if (!callbackName) { + return undefined; + } + const configId = CALLBACK_CONFIG_ALIASES[callbackName] ?? callbackName; + return callbackConfigs.find((config) => config.id === configId); +}; + // Shared helper function to get dynamic params for a callback const getDynamicParamsForCallback = ( callbackName: string | null, @@ -231,7 +272,7 @@ const getDynamicParamsForCallback = ( return fallbackVariables ? Object.keys(fallbackVariables) : []; } - const callbackConfig = callbackConfigs.find((config) => config.id === callbackName); + const callbackConfig = findCallbackConfig(callbackConfigs, callbackName); if (callbackConfig?.dynamic_params) { return Object.keys(callbackConfig.dynamic_params); } From 6e7de5fd20836b499faddda3d7c45bf440545a08 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 23:40:35 +0000 Subject: [PATCH 038/187] test(s3): type the prompts-only test helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/integrations/test_s3.py | 4 +- tests/test_litellm/integrations/test_s3_v2.py | 42 +++++++++++-------- 2 files changed, 27 insertions(+), 19 deletions(-) diff --git a/tests/test_litellm/integrations/test_s3.py b/tests/test_litellm/integrations/test_s3.py index ba8d575c1b7..fd677b9dfdf 100644 --- a/tests/test_litellm/integrations/test_s3.py +++ b/tests/test_litellm/integrations/test_s3.py @@ -31,7 +31,7 @@ def _log_event_kwargs(response_id: str = "chatcmpl-test-id") -> dict: def _run_log_event( - callback_params: dict, response_id: str = "chatcmpl-test-id", log_kwargs: dict | None = None + callback_params: dict, response_id: str = "chatcmpl-test-id", log_kwargs: dict[str, object] | None = None ) -> MagicMock: original = litellm.s3_callback_params litellm.s3_callback_params = callback_params @@ -194,7 +194,7 @@ def test_put_object_keeps_the_configured_path_intact_when_only_the_id_has_to_shr assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES -def _uploaded_body(mock_s3_client: MagicMock) -> dict: +def _uploaded_body(mock_s3_client: MagicMock) -> dict[str, object]: return json.loads(mock_s3_client.put_object.call_args.kwargs["Body"]) diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 1179aa7e409..52fbbe40b0e 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -5,6 +5,7 @@ import re import sys import textwrap import uuid +from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager from datetime import datetime from pathlib import Path @@ -2315,26 +2316,28 @@ def test_build_object_url_uses_partition_dns_suffix(region_name: str, expected_u assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url -def _prompts_only_logger(**kwargs) -> S3Logger: +def _prompts_only_logger(s3_log_prompts_only: bool | None = None) -> S3Logger: return S3Logger( s3_bucket_name="test-bucket", s3_aws_access_key_id="test-key", s3_aws_secret_access_key="test-secret", s3_region_name="us-east-1", - **kwargs, + s3_log_prompts_only=s3_log_prompts_only, ) -def _chat_payload() -> dict: - return { - "id": "chatcmpl-prompts-only", - "messages": [{"role": "user", "content": "Reply with exactly the word PINEAPPLE."}], - "response": {"choices": [{"message": {"role": "assistant", "content": "PINEAPPLE"}}]}, - "metadata": {"user_api_key_team_alias": None}, - } +def _chat_payload() -> StandardLoggingPayload: + return StandardLoggingPayload( + id="chatcmpl-prompts-only", + messages=[{"role": "user", "content": "Reply with exactly the word PINEAPPLE."}], + response={"choices": [{"message": {"role": "assistant", "content": "PINEAPPLE"}}]}, + metadata={"user_api_key_team_alias": None}, + ) -async def _queued_body_via_async_upload(logger: S3Logger, log_event) -> dict: +async def _queued_body_via_async_upload( + logger: S3Logger, log_event: Callable[..., Awaitable[None]] +) -> dict[str, object]: payload = _chat_payload() original = copy.deepcopy(payload) await log_event( @@ -2357,13 +2360,18 @@ async def _queued_body_via_async_upload(logger: S3Logger, log_event) -> dict: @pytest.mark.asyncio @pytest.mark.parametrize("event_name", ["async_log_success_event", "async_log_failure_event"]) -async def test_prompts_only_drops_response_but_keeps_messages_in_uploaded_object(monkeypatch, event_name): +async def test_prompts_only_drops_response_but_keeps_messages_in_uploaded_object( + monkeypatch: pytest.MonkeyPatch, event_name: str +): import litellm monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": True}) logger = _prompts_only_logger() - body = await _queued_body_via_async_upload(logger, getattr(logger, event_name)) + log_event: Callable[..., Awaitable[None]] = ( + logger.async_log_success_event if event_name == "async_log_success_event" else logger.async_log_failure_event + ) + body = await _queued_body_via_async_upload(logger, log_event) assert body["messages"] == _chat_payload()["messages"] assert body["response"] is None @@ -2371,7 +2379,7 @@ async def test_prompts_only_drops_response_but_keeps_messages_in_uploaded_object @pytest.mark.asyncio -async def test_prompts_only_default_off_keeps_response_in_uploaded_object(monkeypatch): +async def test_prompts_only_default_off_keeps_response_in_uploaded_object(monkeypatch: pytest.MonkeyPatch): import litellm monkeypatch.setattr(litellm, "s3_callback_params", {}) @@ -2385,7 +2393,7 @@ async def test_prompts_only_default_off_keeps_response_in_uploaded_object(monkey @pytest.mark.asyncio -async def test_prompts_only_explicit_false_in_params_beats_env_var(monkeypatch): +async def test_prompts_only_explicit_false_in_params_beats_env_var(monkeypatch: pytest.MonkeyPatch): import litellm monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": False}) @@ -2398,7 +2406,7 @@ async def test_prompts_only_explicit_false_in_params_beats_env_var(monkeypatch): @pytest.mark.asyncio -async def test_prompts_only_env_var_applies_when_param_unset(monkeypatch): +async def test_prompts_only_env_var_applies_when_param_unset(monkeypatch: pytest.MonkeyPatch): import litellm monkeypatch.setattr(litellm, "s3_callback_params", {}) @@ -2412,7 +2420,7 @@ async def test_prompts_only_env_var_applies_when_param_unset(monkeypatch): @respx.mock -def test_prompts_only_constructor_kwarg_applies_to_sync_upload(monkeypatch): +def test_prompts_only_constructor_kwarg_applies_to_sync_upload(monkeypatch: pytest.MonkeyPatch): import litellm monkeypatch.setattr(litellm, "s3_callback_params", {}) @@ -2436,7 +2444,7 @@ def test_prompts_only_constructor_kwarg_applies_to_sync_upload(monkeypatch): @pytest.mark.parametrize("callback_name", ["s3", "s3_v2"]) -def test_prompts_only_toggle_is_exposed_to_admin_ui_for_both_s3_callbacks(callback_name): +def test_prompts_only_toggle_is_exposed_to_admin_ui_for_both_s3_callbacks(callback_name: str): from litellm.integrations.custom_logger import CustomLogger assert "S3_LOG_PROMPTS_ONLY" in CustomLogger.get_callback_env_vars(callback_name) From 251aeea97d4e67b5baf4238d851d365f239b5c2e Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 23:56:31 +0000 Subject: [PATCH 039/187] fix(ui): show the S3 label when editing the s3_v2 callback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/components/settings.test.tsx | 1 + ui/litellm-dashboard/src/components/settings.tsx | 8 ++++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/ui/litellm-dashboard/src/components/settings.test.tsx b/ui/litellm-dashboard/src/components/settings.test.tsx index c24fa00438a..4ba5dd23fd1 100644 --- a/ui/litellm-dashboard/src/components/settings.test.tsx +++ b/ui/litellm-dashboard/src/components/settings.test.tsx @@ -393,6 +393,7 @@ describe("Settings", () => { const promptsOnlySwitch = await screen.findByRole("switch", { name: "Log Prompts Only" }); expect(promptsOnlySwitch).not.toBeChecked(); + expect(within(screen.getByRole("dialog")).getByRole("combobox", { name: "Callback" })).toHaveValue("S3"); await user.click(promptsOnlySwitch); await user.click(within(screen.getByRole("dialog")).getByRole("button", { name: "Save Changes" })); diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index c0daeac3b72..9247f22ec28 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -178,7 +178,7 @@ export const CallbackSelector: React.FC = ({ }) => { const { control } = useFormContext(); const inputId = React.useId(); - const selectedConfig = callbackConfigs.find((config) => config.id === selectedCallback) ?? null; + const selectedConfig = findCallbackConfig(callbackConfigs, selectedCallback) ?? null; return ( ; } -const findCallbackConfig = ( - callbackConfigs: readonly CallbackConfigWithParams[], +const findCallbackConfig = ( + callbackConfigs: readonly T[], callbackName: string | null, -): CallbackConfigWithParams | undefined => { +): T | undefined => { if (!callbackName) { return undefined; } From 9541b0734b5aee13bd386a23068504d438424247 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 00:35:45 +0000 Subject: [PATCH 040/187] refactor(keys): drop status helper docstrings and test /key/list status through the endpoint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 16 ++----- .../test_key_management_endpoints.py | 46 +++++++++++++------ 2 files changed, 37 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7a363729541..802a7c3e469 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -6241,7 +6241,6 @@ class _KeyStatusSource(BaseModel): def _derive_key_status(row: Mapping[str, object], now: datetime) -> KeyStatus: - """Status of a live key row; mirrors the partition `_build_status_where_clause` applies at query time.""" source: Final = _KeyStatusSource.model_validate(row) if source.blocked is True: return "revoked" @@ -6651,16 +6650,11 @@ def _not_blocked_where_clause() -> dict[str, object]: def _build_status_where_clause(status_filter: str | None, now: datetime) -> dict[str, object] | None: - """Live-table clause for a status filter; None when the status needs no clause (deleted rows live elsewhere).""" - match status_filter: - case "revoked": - return {"blocked": True} - case "expired": - return {"AND": [_not_blocked_where_clause(), _build_expires_where_clause("expired", now)]} - case "active": - return {"AND": [_not_blocked_where_clause(), _build_expires_where_clause("active", now)]} - case _: - return None + if status_filter == "revoked": + return {"blocked": True} + if status_filter in ("expired", "active"): + return {"AND": [_not_blocked_where_clause(), _build_expires_where_clause(status_filter, now)]} + return None def _build_key_search_where(search: str) -> KeySearchWhere: diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 852b8632846..a8ff860c7c9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6020,22 +6020,46 @@ async def test_list_keys_with_invalid_status(): @pytest.mark.asyncio @pytest.mark.parametrize("status_filter", ["active", "expired", "revoked"]) async def test_list_keys_accepts_live_status_filters(monkeypatch, status_filter): - """LIT-1650: /key/list used to 400 on every status but "deleted"; the live statuses reach the helper.""" from unittest.mock import Mock - from litellm.proxy.management_endpoints import key_management_endpoints + from litellm.proxy.management_endpoints.key_management_endpoints import list_keys - helper = AsyncMock(return_value={"keys": [], "total_count": 0, "current_page": 1, "total_pages": 0}) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock()) - monkeypatch.setattr(key_management_endpoints, "_list_key_helper", helper) - await key_management_endpoints.list_keys( + live_row = MagicMock() + live_row.model_dump.return_value = {"token": "hashed_live_token", "object_permission_id": None} + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[live_row]) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + response = await list_keys( request=Mock(), user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + page=1, + size=10, + user_id=None, + team_id=None, + organization_id=None, + key_hash=None, + key_alias=None, + search=None, + return_full_object=False, + include_team_keys=False, + include_created_by_keys=False, + sort_by=None, + sort_order="desc", + expand=None, status=status_filter, + project_id=None, + access_group_id=None, + agent_id=None, + substring_matching=False, + expires=None, ) - assert helper.await_args is not None - assert helper.await_args.kwargs["status"] == status_filter + assert response["keys"] == ["hashed_live_token"] + assert response["total_count"] == 1 + mock_prisma_client.db.litellm_deletedverificationtoken.find_many.assert_not_called() def _status_filter_where(status_filter: str | None) -> Mapping[str, object]: @@ -6054,7 +6078,6 @@ def _status_filter_where(status_filter: str | None) -> Mapping[str, object]: def test_build_key_filter_conditions_status_filter_partitions_live_keys(): - """LIT-1650: active, expired and revoked are disjoint predicates over blocked + expires on the live table.""" not_blocked = {"OR": [{"blocked": None}, {"blocked": False}]} revoked_where = _status_filter_where("revoked") @@ -6072,13 +6095,11 @@ def test_build_key_filter_conditions_status_filter_partitions_live_keys(): def test_build_key_filter_conditions_deleted_status_adds_no_live_clause(): - """Deleted rows live in the archive table, so the status must not narrow the live-table query.""" assert _status_filter_where("deleted") == _status_filter_where(None) @pytest.mark.asyncio async def test_list_key_helper_revoked_status_filters_live_table_on_blocked(): - """LIT-1650: status="revoked" stays on the live table and narrows it to blocked keys.""" mock_prisma_client = AsyncMock() mock_find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many @@ -6123,7 +6144,6 @@ def _archived_key_row(token: str, user_id: str) -> MagicMock: @pytest.mark.asyncio async def test_info_key_fn_serves_deleted_key_from_archive(monkeypatch): - """LIT-1650: /key/info falls back to LiteLLM_DeletedVerificationToken and reports status="deleted".""" from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn hashed = "hashed_deleted_token" @@ -6153,7 +6173,6 @@ async def test_info_key_fn_serves_deleted_key_from_archive(monkeypatch): @pytest.mark.asyncio async def test_info_key_fn_archived_key_keeps_owner_authorization(monkeypatch): - """An archived key is still scoped: a different internal user gets 403, the owner gets the record.""" from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn hashed = "hashed_deleted_token" @@ -6210,7 +6229,6 @@ async def test_info_key_fn_unknown_key_still_404s(monkeypatch): ], ) async def test_info_key_fn_reports_live_key_status(monkeypatch, blocked, expires, expected_status): - """LIT-1650: live keys carry the same status vocabulary /key/list filters on.""" from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn mock_prisma_client = AsyncMock() From 260ff5f491e08d0c357a39c8b99a286403da8255 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 23:39:29 +0000 Subject: [PATCH 041/187] feat(team): team-level model_max_budget with key-level overrides A team can now carry a per-model budget map that every key on the team inherits. A key's own model_max_budget entry for the same model takes precedence, so it is gated on and billed to the key alone. Backend: NewTeamRequest/UpdateTeamRequest accept model_max_budget (validated like the key-level field, enterprise gated); the value is hydrated onto UserAPIKeyAuth via the token view, TeamGrants and the carried budget state; _check_team_model_budget enforces it in the centralized common checks; the limiter meters spend under team_model_spend::: and skips the team counter when the key overrides; /team/update lets only a proxy admin raise, re-window or drop a cap; /team/info exposes usage. The Anthropic context-management compaction summary subrequest runs the same team gate. Both fallback token-view SQL definitions project the column. UI: team create and edit forms reuse the key-level ModelMaxBudgetEditor, premium gated, sending {} to clear and omitting unchanged fields. A key entry overrides the team cap only when it spend-gates the model (non-negative max_budget); a row that only carries tpm/rpm limits or a negative cap leaves the team cap in force. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../context_management/editors/compact.py | 27 +- litellm/proxy/_types.py | 16 ++ litellm/proxy/auth/team_grants.py | 2 + litellm/proxy/auth/user_api_key_auth.py | 47 +++ litellm/proxy/db/create_views.py | 1 + .../proxy/hooks/model_max_budget_limiter.py | 58 +++- litellm/proxy/litellm_pre_call_utils.py | 1 + .../management_endpoints/common_utils.py | 51 ++++ .../management_endpoints/team_endpoints.py | 77 ++++- .../pass_through_endpoints.py | 1 + .../spend_tracking/carried_budget_state.py | 1 + litellm/proxy/utils.py | 2 + ...test_unit_test_max_model_budget_limiter.py | 267 ++++++++++++++++++ .../context_management/test_compact.py | 74 +++++ .../proxy/auth/test_team_grants.py | 2 + .../proxy/auth/test_user_api_key_auth.py | 68 +++++ .../proxy/db/test_create_views.py | 1 + .../management_endpoints/test_common_utils.py | 52 ++++ .../test_team_endpoints.py | 229 +++++++++++++++ .../test_carried_budget_state.py | 13 + .../test_prisma_client_get_data.py | 11 +- .../src/components/Teams.test.tsx | 31 ++ ui/litellm-dashboard/src/components/Teams.tsx | 15 + .../key_team_helpers/ModelMaxBudgetEditor.tsx | 1 + .../src/components/team/TeamInfo.test.tsx | 93 ++++++ .../src/components/team/TeamInfo.tsx | 36 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 27 ++ 27 files changed, 1192 insertions(+), 12 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index fb6a1c40253..ecaf8f2e7e1 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -78,6 +78,7 @@ _PROPAGATED_METADATA_KEYS: Final = ( "user_api_key_end_user_id", "user_api_end_user_max_budget", "user_api_key_model_max_budget", + "user_api_key_team_model_max_budget", "user_api_key_user_model_max_budget", "user_api_key_end_user_model_max_budget", "litellm_call_id", @@ -395,9 +396,9 @@ async def _check_summary_model_budget( ``user_api_key_auth`` runs for the client-requested model. Returns True outside the proxy or when no per-model budget is configured. - All three scopes are checked because the summary's spend is charged to all - three: this file propagates the key, user and end-user budgets into the - subrequest's metadata, so enforcing only two of them would let compaction + Every scope is checked because the summary's spend is charged to every + scope: this file propagates the key, team, user and end-user budgets into the + subrequest's metadata, so skipping one of them would let compaction increment a counter it can never be refused by. """ if user_api_key_auth is None: @@ -444,6 +445,26 @@ async def _check_summary_model_budget( ) return False + team_model_max_budget: Final = user_api_key_auth.team_model_max_budget + team_id: Final = user_api_key_auth.team_id + if isinstance(team_model_max_budget, dict) and team_model_max_budget and team_id is not None: + try: + await model_max_budget_limiter.is_team_within_model_budget( + team_id=team_id, + team_model_max_budget=team_model_max_budget, + key_model_max_budget=model_max_budget if isinstance(model_max_budget, dict) else None, + model=summary_model, + ) + except litellm.BudgetExceededError: + return False + except Exception as e: # noqa: BLE001 # a budget gate denies on any failure, as the other scopes do + verbose_logger.warning( + "compact_20260112: unexpected error during team model-budget check for summary_model=%s; denying: %s", + summary_model, + e, + ) + return False + end_user_model_max_budget: Final[dict[str, object] | None] = getattr( user_api_key_auth, "end_user_model_max_budget", None ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ad55fa5d2be..72852c7c20c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2002,6 +2002,13 @@ RouterSettingsDict = Annotated[ class NewTeamRequest(TeamBase): router_settings: RouterSettingsDict | None = None model_aliases: dict | None = None + model_max_budget: GenericBudgetConfigType | None = Field( + default=None, + description=( + "Max budget per model for every key on the team, overridable per key " + "(e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}})" + ), + ) tags: list | None = None guardrails: list[str] | None = None policies: list[str] | None = None @@ -2103,6 +2110,13 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): access_group_ids: list[str] | None = None budget_limits: list[BudgetLimitEntry] | None = None # multiple concurrent budget windows default_team_member_models: list[str] | None = None # default allowed_models seeded onto new team members + model_max_budget: GenericBudgetConfigType | None = Field( + default=None, + description=( + "Max budget per model for every key on the team, overridable per key " + "(e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}})" + ), + ) class PatchTeamRequest(UpdateTeamRequest): @@ -3030,6 +3044,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): team_tpd_limit: int | None = None team_max_budget: float | None = None team_soft_budget: float | None = None + team_model_max_budget: dict[str, object] | None = None team_models: list = [] team_blocked: bool = False soft_budget: float | None = None @@ -4444,6 +4459,7 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): # Parent org's model ceiling, reported only to callers who can manage the team. # None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling. organization_models: list[str] | None = None + model_max_budget_usage: Mapping[str, Mapping[str, object]] | None = None class TeamInfoResponseObject(TypedDict): diff --git a/litellm/proxy/auth/team_grants.py b/litellm/proxy/auth/team_grants.py index 0421659c331..2029ee342ae 100644 --- a/litellm/proxy/auth/team_grants.py +++ b/litellm/proxy/auth/team_grants.py @@ -59,6 +59,7 @@ class TeamGrants(TypedDict, total=False): team_tpd_limit: ReadOnly[int | None] team_max_budget: ReadOnly[float | None] team_soft_budget: ReadOnly[float | None] + team_model_max_budget: ReadOnly[dict[str, object] | None] team_spend: ReadOnly[float | None] team_models: ReadOnly[Sequence[str]] team_blocked: ReadOnly[bool] @@ -101,6 +102,7 @@ def team_grants( team_tpd_limit=team_object.tpd_limit, team_max_budget=team_object.max_budget, team_soft_budget=team_object.soft_budget, + team_model_max_budget=team_object.model_max_budget, team_spend=team_object.spend, team_models=tuple(team_object.models), team_blocked=team_object.blocked, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5958a68f975..0a5aa9ea793 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -264,6 +264,16 @@ class _UserModelBudgetLimiter(Protocol): ) -> bool: ... +class _TeamModelBudgetLimiter(Protocol): + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, object], + key_model_max_budget: Mapping[str, object] | None, + model: str, + ) -> bool: ... + + class _TokenTeamModels(Protocol): @property def team_models(self) -> list[str]: ... @@ -334,6 +344,25 @@ async def _check_user_model_budget( ) +async def _check_team_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: _TeamModelBudgetLimiter, + models: list[str], +) -> None: + """Enforce the team's `model_max_budget` for every requested model the key does not override.""" + team_model_max_budget: Final = valid_token.team_model_max_budget + if valid_token.team_id is None or not team_model_max_budget: + return + key_model_max_budget: Final[Mapping[str, object] | None] = valid_token.model_max_budget + for model_name in models: + await model_max_budget_limiter.is_team_within_model_budget( + team_id=valid_token.team_id, + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model=model_name, + ) + + async def _check_key_model_budget_with_fallback( valid_token: UserAPIKeyAuth, model_max_budget_limiter: _KeyModelBudgetLimiter, @@ -2369,6 +2398,7 @@ async def _user_api_key_auth_builder( team_id=valid_token.team_id, max_budget=valid_token.team_max_budget, soft_budget=valid_token.team_soft_budget, + model_max_budget=valid_token.team_model_max_budget, spend=valid_token.team_spend, tpm_limit=valid_token.team_tpm_limit, rpm_limit=valid_token.team_rpm_limit, @@ -2523,6 +2553,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached team_id=valid_token.team_id, max_budget=valid_token.team_max_budget, soft_budget=valid_token.team_soft_budget, + model_max_budget=valid_token.team_model_max_budget, spend=valid_token.team_spend, tpm_limit=valid_token.team_tpm_limit, rpm_limit=valid_token.team_rpm_limit, @@ -2592,6 +2623,7 @@ async def _run_centralized_common_checks( litellm_proxy_admin_name, llm_router, master_key, + model_max_budget_limiter, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -2864,6 +2896,21 @@ async def _run_centralized_common_checks( finally: release_spend_counter_batch() + if not skip_budget_checks: + await _check_team_model_budget( + valid_token=user_api_key_auth_obj, + model_max_budget_limiter=model_max_budget_limiter, + models=_get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=user_api_key_auth_obj.team_id, + ) + ), + ) + await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, request=request, diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index d3f3de730ab..f7131091c0b 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -78,6 +78,7 @@ async def create_missing_views(db: SupportsRawQueries) -> None: v.*, t.spend AS team_spend, t.max_budget AS team_max_budget, + t.model_max_budget AS team_model_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.tpd_limit AS team_tpd_limit, diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index efaaab277a9..bbfc7325f40 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -19,12 +19,14 @@ from litellm.types.utils import BudgetConfig, StandardLoggingPayload VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend" END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend" USER_SPEND_CACHE_KEY_PREFIX: Final = "user_model_spend" +TEAM_SPEND_CACHE_KEY_PREFIX: Final = "team_model_spend" _SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType( { Litellm_EntityType.KEY: VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, Litellm_EntityType.USER: USER_SPEND_CACHE_KEY_PREFIX, Litellm_EntityType.END_USER: END_USER_SPEND_CACHE_KEY_PREFIX, + Litellm_EntityType.TEAM: TEAM_SPEND_CACHE_KEY_PREFIX, } ) @@ -37,6 +39,7 @@ _BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType( Litellm_EntityType.KEY: "virtual_key_budget_start_time", Litellm_EntityType.USER: "user_model_budget_start_time", Litellm_EntityType.END_USER: "end_user_budget_start_time", + Litellm_EntityType.TEAM: "team_model_budget_start_time", } ) @@ -139,6 +142,18 @@ def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> return None +def team_model_budget_applies(model: str, key_model_max_budget: Mapping[str, object] | None) -> bool: + """A key entry that spend-gates `model` overrides the team cap: it is then gated on and billed to the key alone.""" + if not key_model_max_budget: + return True + resolved: Final = resolve_model_budget(model=model, model_max_budget=key_model_max_budget) + return resolved is None or not _spend_gated(resolved.budget_config) + + +def _spend_gated(budget_config: BudgetConfig) -> bool: + return budget_config.max_budget is not None and budget_config.max_budget >= 0 + + def _budget_model_candidates(model: str) -> tuple[str, ...]: """Names a budget may be configured under for a request on `model`, most specific first. @@ -346,6 +361,30 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): exceeded_message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", ) + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, object], + key_model_max_budget: Mapping[str, object] | None, + model: str, + ) -> bool: + """ + Check if the team is within the model budget, unless the key's own + `model_max_budget` overrides it for `model` + + Raises: + BudgetExceededError: If the team has exceeded the model budget + """ + if not team_model_budget_applies(model=model, key_model_max_budget=key_model_max_budget): + return True + return await self._is_entity_within_model_budget( + entity_type=Litellm_EntityType.TEAM, + entity_id=team_id, + model_max_budget=team_model_max_budget, + model=model, + exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}", + ) + async def _is_entity_within_model_budget( self, entity_type: Litellm_EntityType, @@ -456,11 +495,26 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) + key_model_max_budget: Final = _metadata.get("user_api_key_model_max_budget") entity_budgets: Final = ( ( Litellm_EntityType.KEY, payload_metadata.get("user_api_key_hash"), - _metadata.get("user_api_key_model_max_budget"), + key_model_max_budget, + ), + ( + Litellm_EntityType.TEAM, + payload_metadata.get("user_api_key_team_id"), + ( + _metadata.get("user_api_key_team_model_max_budget") + if team_model_budget_applies( + model=model, + key_model_max_budget=( + key_model_max_budget if isinstance(key_model_max_budget, Mapping) else None + ), + ) + else None + ), ), ( Litellm_EntityType.USER, @@ -478,7 +532,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): if not resolved_budgets: verbose_proxy_logger.debug( "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: " - "no key, user or end-user model_max_budget covers model=%s", + "no key, team, user or end-user model_max_budget covers model=%s", model, ) return diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 563db811edc..93374eb099b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2287,6 +2287,7 @@ async def add_litellm_data_to_request( # Team spend, budget - used by prometheus.py data[_metadata_variable_name]["user_api_key_team_max_budget"] = user_api_key_dict.team_max_budget data[_metadata_variable_name]["user_api_key_team_spend"] = user_api_key_dict.team_spend + data[_metadata_variable_name]["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget data[_metadata_variable_name]["user_api_key_request_route"] = user_api_key_dict.request_route # API Key spend, budget - used by prometheus.py diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 98155ad6839..973311608ed 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -55,6 +55,7 @@ def validate_budget_duration(budget_duration: str | None, status_code: int = 400 from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import ( + CommonProxyErrors, KeyRequestBase, LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, @@ -73,12 +74,62 @@ from litellm.proxy._types import ( # noqa: F401 re-exported from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.utils import _premium_user_check from litellm.repositories.team_repository import TeamRepository +from litellm.types.utils import BudgetConfig if TYPE_CHECKING: from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest from litellm.proxy.utils import PrismaClient, ProxyLogging +def validate_team_model_max_budget( + model_max_budget: Mapping[str, BudgetConfig] | None, + premium_user: bool, +) -> None: + """Reject a team `model_max_budget` the limiter could not enforce (no duration, bad cap, tpm/rpm limits).""" + if not model_max_budget: + return + if premium_user is not True: + raise HTTPException( + status_code=403, + detail={ + "error": f"Setting model_max_budget on a team is an enterprise feature. {CommonProxyErrors.not_premium_user.value}" + }, + ) + for model_name, budget_config in model_max_budget.items(): + if not model_name.strip(): + raise HTTPException( + status_code=400, + detail={"error": "model_max_budget keys must be non-empty model names"}, + ) + max_budget = budget_config.max_budget + if max_budget is None or not math.isfinite(max_budget) or max_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"model_max_budget[{model_name!r}].max_budget must be a non-negative finite number. " + f"Received: {max_budget}" + ) + }, + ) + if budget_config.budget_duration is None: + raise HTTPException( + status_code=400, + detail={"error": f"model_max_budget[{model_name!r}] requires a budget_duration, e.g. '1d' or '30d'"}, + ) + validate_budget_duration(budget_config.budget_duration) + if budget_config.tpm_limit is not None or budget_config.rpm_limit is not None: + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"model_max_budget[{model_name!r}] tpm_limit/rpm_limit are not enforced on a team; " + "set per-model rate limits on the key instead" + ) + }, + ) + + def require_caller_user_id_for_non_admin( user_api_key_dict: UserAPIKeyAuth, ) -> str: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e719d6d761a..6fb1ef5ec93 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protoc import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, JsonValue, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -38,6 +38,7 @@ from litellm.proxy._types import ( DeleteTeamRequest, LiteLLM_AuditLogs, LiteLLM_DeletedTeamTable, + Litellm_EntityType, LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_ModelTable, @@ -95,6 +96,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity_aggregated, ) @@ -108,6 +110,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _upsert_budget_and_membership, _user_has_admin_view, validate_budget_duration, + validate_team_model_max_budget, ) from litellm.proxy.management_endpoints.organization_endpoints import ( add_member_to_organization, @@ -177,6 +180,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( TeamUserSpendRow, UpdateTeamMemberPermissionsRequest, ) +from litellm.types.utils import BudgetConfig if TYPE_CHECKING: from prisma import Prisma @@ -1170,6 +1174,56 @@ def _check_team_budget_update_authority( ) +def _existing_model_cap(raw_budget_config: object) -> BudgetConfig | None: + try: + return BudgetConfig.model_validate(raw_budget_config) + except ValidationError: + return None + + +def _check_team_model_budget_update_authority( + data: UpdateTeamRequest, + user_api_key_dict: UserAPIKeyAuth, + existing_model_max_budget: Mapping[str, object] | None, +) -> None: + """Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap.""" + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + if "model_max_budget" not in data.model_fields_set or not existing_model_max_budget: + return + requested: Final[Mapping[str, BudgetConfig]] = data.model_max_budget or {} + for model_name, raw_existing in existing_model_max_budget.items(): + existing = _existing_model_cap(raw_existing) + if existing is None or existing.max_budget is None: + continue + proposed = requested.get(model_name) + if proposed is None: + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"Only a proxy admin can remove a team's model_max_budget for {model_name!r}. " + f"Current max_budget={existing.max_budget}." + ) + }, + ) + if ( + proposed.max_budget is None + or proposed.max_budget > existing.max_budget + or proposed.budget_duration != existing.budget_duration + ): + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"Only a proxy admin can raise a team's model_max_budget for {model_name!r} or change its " + f"budget_duration. Current max_budget={existing.max_budget} per {existing.budget_duration}, " + f"requested={proposed.max_budget} per {proposed.budget_duration}." + ) + }, + ) + + def _should_auto_add_team_creator( user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], @@ -1230,6 +1284,7 @@ async def new_team( - prompts: Optional[List[str]] - List of prompts that the team is allowed to use. - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -1291,6 +1346,7 @@ async def new_team( general_settings, litellm_proxy_admin_name, llm_router, + premium_user, prisma_client, user_api_key_cache, ) @@ -1321,6 +1377,7 @@ async def new_team( validate_budget_duration(data.budget_duration) validate_budget_duration(data.team_member_budget_duration) + validate_team_model_max_budget(model_max_budget=data.model_max_budget, premium_user=premium_user) if data.soft_budget is not None: if data.max_budget is not None: @@ -1980,6 +2037,7 @@ async def update_team( - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -2031,6 +2089,7 @@ async def update_team( from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, llm_router, + premium_user, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -2069,6 +2128,7 @@ async def update_team( validate_budget_duration(data.budget_duration) validate_budget_duration(data.team_member_budget_duration) + validate_team_model_max_budget(model_max_budget=data.model_max_budget, premium_user=premium_user) existing_team_row = await _raw_team_db(TeamRepository(prisma_client)).find_unique( where={"team_id": data.team_id} @@ -2204,8 +2264,15 @@ async def update_team( user_api_key_dict=user_api_key_dict, existing_team_max_budget=existing_team_row.max_budget, ) + _check_team_model_budget_update_authority( + data=data, + user_api_key_dict=user_api_key_dict, + existing_model_max_budget=existing_team_row.model_max_budget, + ) updated_kv = data.json(exclude_unset=True) + if "model_max_budget" in updated_kv and updated_kv["model_max_budget"] is None: + updated_kv["model_max_budget"] = {} # Drop server-owned metadata keys from caller input so they can only # be written by the same code path that creates the underlying rows. @@ -4473,7 +4540,7 @@ async def team_info( ``` """ from litellm.proxy._types import TeamInfoResponseObjectTeamTable - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client try: if prisma_client is None: @@ -4573,6 +4640,12 @@ async def team_info( update={ # mutable-ok: pydantic update payload "members_with_roles": hydrated_members, "organization_models": organization_models, + "model_max_budget_usage": await build_model_max_budget_usage( + entity_type=Litellm_EntityType.TEAM, + entity_id=team_id, + model_max_budget=resolved_team_info.model_max_budget, + cache=model_max_budget_limiter.dual_cache, + ), } ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 686544d352c..de6f9cb7647 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -609,6 +609,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): # merely shares the name. if not request_dispatched_to_pass_through_endpoint(request): _metadata["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget + _metadata["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget _metadata["user_api_key_user_model_max_budget"] = user_api_key_dict.user_model_max_budget _metadata["user_api_key_end_user_model_max_budget"] = user_api_key_dict.end_user_model_max_budget _metadata.update( diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index efd3a78d211..da8bf60ebda 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -25,6 +25,7 @@ def carry_team_and_user_budget_state( budget_reset_at=team_object.budget_reset_at, max_budget=team_object.max_budget, ) + valid_token.team_model_max_budget = team_object.model_max_budget # rebind-ok: caller keeps this object if user_object is not None: valid_token.user_budget_snapshot = UserBudgetSnapshot( # rebind-ok: same object the caller keeps using budget_reset_at=user_object.budget_reset_at, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 479bd0a55af..b77f25389c6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4319,6 +4319,7 @@ class PrismaClient: v.*, t.spend AS team_spend, t.max_budget AS team_max_budget, + t.model_max_budget AS team_model_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.tpd_limit AS team_tpd_limit @@ -4758,6 +4759,7 @@ class PrismaClient: t.spend AS team_spend, t.max_budget AS team_max_budget, t.soft_budget AS team_soft_budget, + t.model_max_budget AS team_model_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.tpd_limit AS team_tpd_limit, diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 096efc33aaf..efe41e1da9a 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -587,6 +587,8 @@ def _success_kwargs( response_cost=0.5, key_hash=None, key_model_max_budget=None, + team_id=None, + team_model_max_budget=None, user_id=None, user_model_max_budget=None, end_user_id=None, @@ -600,6 +602,7 @@ def _success_kwargs( "end_user": end_user_id, "metadata": { "user_api_key_hash": key_hash, + "user_api_key_team_id": team_id, "user_api_key_user_id": user_id, "user_api_key_end_user_id": end_user_id, }, @@ -607,6 +610,7 @@ def _success_kwargs( "litellm_params": { "metadata": { "user_api_key_model_max_budget": key_model_max_budget, + "user_api_key_team_model_max_budget": team_model_max_budget, "user_api_key_user_model_max_budget": user_model_max_budget, "user_api_key_end_user_model_max_budget": end_user_model_max_budget, }, @@ -1417,3 +1421,266 @@ async def test_spend_logged_on_one_replica_is_enforced_and_reported_on_another() replica_c = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) with pytest.raises(litellm.BudgetExceededError): await replica_c.is_key_within_model_budget(user_api_key, "gpt-4") + + +def _log_success(limiter, **kwargs): + return limiter.async_log_success_event( + _success_kwargs(**kwargs), response_obj=None, start_time=None, end_time=None + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_model", + ["gpt-4", "openai/gpt-4"], + ids=["bare_model", "provider_prefixed_model"], +) +async def test_team_model_budget_is_shared_by_every_key_without_an_override(request_model): + """ + Two keys on the same team, neither carrying a matching key-level entry, + charge one team counter and are both refused once it is spent. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + check = lambda: limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model=request_model, + ) + + assert await check() is True + await _log_success( + limiter, + model_group=request_model, + response_cost=0.6, + key_hash="vk-a", + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + assert await check() is True + await _log_success( + limiter, + model_group=request_model, + response_cost=0.6, + key_hash="vk-b", + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == pytest.approx(1.2) + with pytest.raises(litellm.BudgetExceededError) as exc: + await check() + assert exc.value.entity_type == Litellm_EntityType.TEAM.value + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.TEAM, + entity_id="team-1", + model_max_budget=team_model_max_budget, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": pytest.approx(1.2), "budget_limit": 1.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_key_override_replaces_the_team_cap_for_that_model(): + """ + A key with its own entry for the model is gated on the key counter alone: + the exhausted team counter does not block it, and its spend never lands on + the team counter. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"gpt-4": {"budget_limit": 5.0, "time_period": "1d"}} + await dual_cache.async_set_cache(key="team_model_spend:team-1:gpt-4:1d", value=9.0) + + assert ( + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="openai/gpt-4", + ) + is True + ) + + await _log_success( + limiter, + model_group="openai/gpt-4", + response_cost=2.0, + key_hash="vk-override", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 9.0 + assert await dual_cache.async_get_cache(key="virtual_key_spend:vk-override:gpt-4:1d") == 2.0 + + +@pytest.mark.asyncio +async def test_key_entry_for_another_model_does_not_lift_the_team_cap(): + """A key override only covers the model it names; other models stay on the team counter.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"claude-3": {"budget_limit": 5.0, "time_period": "1d"}} + + await _log_success( + limiter, + model_group="gpt-4", + response_cost=1.5, + key_hash="vk-other", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="gpt-4", + ) + + +@pytest.mark.asyncio +async def test_team_budget_leaves_unconfigured_models_alone(): + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 0.0, "time_period": "1d"}} + + assert ( + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model="claude-3", + ) + is True + ) + with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment: + await _log_success( + limiter, + model_group="claude-3", + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + mock_increment.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_counters_are_isolated_by_team_model_and_window(): + """Same model on two teams, and two models with different windows on one team, never share a counter.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = { + "gpt-4": {"budget_limit": 10.0, "time_period": "1d"}, + "claude-3": {"budget_limit": 10.0, "time_period": "30d"}, + } + + for team_id, model in (("team-1", "gpt-4"), ("team-2", "gpt-4"), ("team-1", "claude-3")): + await _log_success( + limiter, + model_group=model, + response_cost=1.0, + team_id=team_id, + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.0 + assert await dual_cache.async_get_cache(key="team_model_spend:team-2:gpt-4:1d") == 1.0 + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:claude-3:30d") == 1.0 + assert await dual_cache.async_get_cache(key="team_model_budget_start_time:team-1:claude-3:30d") is not None + + +@pytest.mark.asyncio +async def test_malformed_team_entry_is_skipped_and_its_sibling_still_enforced(): + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + team_model_max_budget = { + "gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}, + "claude-3": {"budget_limit": 0.0, "time_period": "1d"}, + } + + assert ( + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model="gpt-4", + ) + is True + ) + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model="claude-3", + ) + + +@pytest.mark.asyncio +async def test_malformed_key_entry_does_not_count_as_an_override(): + """A key entry the limiter cannot enforce must not also switch the team cap off.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}} + + await _log_success( + limiter, + model_group="gpt-4", + response_cost=1.5, + key_hash="vk-bad", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="gpt-4", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_entry", + [ + {"time_period": "1d", "tpm_limit": 100}, + {"time_period": "1d", "rpm_limit": 10}, + {"budget_limit": -1.0, "time_period": "1d"}, + ], +) +async def test_key_entry_without_a_spend_cap_does_not_lift_the_team_cap(key_entry): + """A key row that only rate-limits the model, or has no enforceable cap, leaves the team cap in force.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"gpt-4": key_entry} + + await _log_success( + limiter, + model_group="openai/gpt-4", + response_cost=1.5, + key_hash="vk-rate-limited", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="openai/gpt-4", + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 7660a8649b5..fc5d807bc23 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -1200,6 +1200,7 @@ def _fake_user_api_key_auth( team_models=None, team_id=None, model_max_budget=None, + team_model_max_budget=None, end_user_model_max_budget=None, end_user_id=None, user_model_max_budget=None, @@ -1220,6 +1221,7 @@ def _fake_user_api_key_auth( auth.team_id = team_id auth.team_model_aliases = None auth.model_max_budget = model_max_budget + auth.team_model_max_budget = team_model_max_budget auth.end_user_model_max_budget = end_user_model_max_budget auth.end_user_id = end_user_id auth.user_model_max_budget = user_model_max_budget @@ -1860,6 +1862,78 @@ async def test_summary_model_rate_limit_skipped_for_legacy_limiter(): assert not result.applied_edits[0].get("error") +async def test_summary_model_denied_when_team_over_model_budget(): + """The team per-model budget gates the summary subrequest, whose spend is + charged to the team counter via the propagated `user_api_key_team_model_max_budget`. + The key's own `model_max_budget` is handed to the limiter so a key-level + override keeps taking precedence over the team cap here as it does in auth.""" + import litellm + + messages = _simple_messages() + mock_call = AsyncMock(return_value=_make_mock_response("x")) + key_budget = {"claude-opus-4-8": {"budget_limit": 1}} + team_budget = {"claude-haiku-4-5": {"budget_limit": 5, "time_period": "1d"}} + + auth = _fake_user_api_key_auth( + key_models=["all-proxy-models"], + model_max_budget=key_budget, + team_model_max_budget=team_budget, + team_id="team-over-budget", + token="hashed-token", + ) + + limiter = MagicMock() + limiter.is_key_within_model_budget = AsyncMock(return_value=True) + limiter.is_team_within_model_budget = AsyncMock( + side_effect=litellm.BudgetExceededError( + message="over budget", current_cost=10, max_budget=5 + ) + ) + + with ( + patch( # test-quality-ok: apply_compact_20260112 reads the summary model setting as a module global, no seam + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting", + return_value="claude-haiku-4-5", + ), + patch("litellm.token_counter", return_value=200_000), # test-quality-ok: forces the over-threshold branch + patch( # test-quality-ok: the summary call is the observable that must NOT happen when the team is over budget + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model", + mock_call, + ), + patch( # test-quality-ok: the limiter is a proxy_server module global the editor imports, no injection seam + "litellm.proxy.proxy_server.model_max_budget_limiter", limiter + ), + ): + result = await apply_compact_20260112( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec=_EDIT_SPEC_DEFAULT, + user_api_key_auth=auth, + ) + + mock_call.assert_not_awaited() + assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded" + limiter.is_team_within_model_budget.assert_awaited_once_with( + team_id="team-over-budget", + team_model_max_budget=team_budget, + key_model_max_budget=key_budget, + model="claude-haiku-4-5", + ) + import inspect + + from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, + ) + + real_params = inspect.signature( + _PROXY_VirtualKeyModelMaxBudgetLimiter.is_team_within_model_budget + ).parameters + for kwarg in ("team_id", "team_model_max_budget", "key_model_max_budget", "model"): + assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter does not accept" + + async def test_scoped_budget_metadata_propagated_to_summary_call(): """The end-user/project scope identifiers and the end-user budget the post-call spend and rate-limit hooks key on are forwarded to the summary subrequest, and diff --git a/tests/test_litellm/proxy/auth/test_team_grants.py b/tests/test_litellm/proxy/auth/test_team_grants.py index 7b6717f804f..f74531beaf0 100644 --- a/tests/test_litellm/proxy/auth/test_team_grants.py +++ b/tests/test_litellm/proxy/auth/test_team_grants.py @@ -31,6 +31,7 @@ def _full_team(model_aliases=ALIASES) -> LiteLLM_TeamTable: max_budget=50.0, soft_budget=25.0, spend=12.5, + model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}, models=["gpt-4o", "gpt-4o-mini"], blocked=True, metadata={"tier": "gold"}, @@ -72,6 +73,7 @@ def test_team_grants_cover_every_team_field_the_key_path_gets(): assert token.team_max_budget == 50.0 assert token.team_soft_budget == 25.0 assert token.team_spend == 12.5 + assert token.team_model_max_budget == {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} assert token.team_models == ["gpt-4o", "gpt-4o-mini"] assert token.team_blocked is True assert token.team_metadata == {"tier": "gold"} diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index bd7ff62ac8b..c55dd966b2b 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4372,6 +4372,74 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t } +class _RecordingTeamModelBudgetLimiter: + def __init__(self): + self.calls = [] + + async def is_team_within_model_budget(self, team_id, team_model_max_budget, key_model_max_budget, model): + self.calls.append((team_id, dict(team_model_max_budget), key_model_max_budget, model)) + return True + + +@pytest.mark.asyncio +async def test_centralized_common_checks_enforces_team_model_max_budget_from_the_resolved_team(): + """The team's model_max_budget is enforced at the single authz gate, off the + team object auth resolved (not the possibly stale token copy), and the key's + own model_max_budget is handed to the limiter so a matching key entry can + override the team cap.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + team_caps = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} + key_caps = {"claude-sonnet-4-6": {"max_budget": 1.0, "budget_duration": "1d"}} + token = UserAPIKeyAuth( + api_key="sk-test", + token="hashed", + team_id="t1", + team_model_max_budget={"gpt-4o": {"max_budget": 999.0, "budget_duration": "30d"}}, + model_max_budget=key_caps, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="team_id:t1", + value=LiteLLM_TeamTableCachedObj(team_id="t1", model_max_budget=team_caps), + ) + limiter = _RecordingTeamModelBudgetLimiter() + attrs = { + **_proxy_attrs_for_centralized_checks(user_custom_auth=None), + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "model_max_budget_limiter": limiter, + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), # test-quality-ok: stubs the sibling check so only the team model-budget gate is under test + patch( # test-quality-ok: stubs the budget reservation so only the team model-budget gate is under test + "litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks", + new_callable=AsyncMock, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert limiter.calls == [("t1", team_caps, key_caps, "gpt-4o")] + + @pytest.mark.asyncio async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): """Existing RPS guarantee: custom-auth deployments without diff --git a/tests/test_litellm/proxy/db/test_create_views.py b/tests/test_litellm/proxy/db/test_create_views.py index ecc6d70123e..54418e10bdf 100644 --- a/tests/test_litellm/proxy/db/test_create_views.py +++ b/tests/test_litellm/proxy/db/test_create_views.py @@ -71,6 +71,7 @@ async def test_create_views_creates_view_on_does_not_exist(): mock_db.execute_raw.assert_called_once() created_sql = mock_db.execute_raw.call_args[0][0] assert 'CREATE VIEW "LiteLLM_VerificationTokenView"' in created_sql + assert "t.model_max_budget AS team_model_max_budget" in created_sql @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index 7352ca0e9ee..2b614632346 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -35,6 +35,7 @@ from litellm.proxy.management_endpoints.common_utils import ( admin_can_invite_user, ) from litellm.proxy.management_endpoints.common_utils import _has_non_empty_value +from litellm.types.utils import BudgetConfig class TestUpdateMetadataFieldsEmptyCollections: @@ -1162,3 +1163,54 @@ async def test_router_weights_validate_current_deployment_scope( assert exc.value.detail == error else: await validation + + +@pytest.mark.parametrize( + "model_max_budget, error", + [ + ({"gpt-4o": BudgetConfig(max_budget=-1.0, budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(max_budget=float("inf"), budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(max_budget=float("nan"), budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(max_budget=5.0)}, "requires a budget_duration"), + ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="fortnight")}, "budget_duration"), + ({" ": BudgetConfig(max_budget=5.0, budget_duration="1d")}, "non-empty model names"), + ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="1d", tpm_limit=1000)}, "not enforced on a team"), + ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="1d", rpm_limit=10)}, "not enforced on a team"), + ], + ids=["negative", "inf", "nan", "no_cap", "no_duration", "bad_duration", "blank_model", "tpm_limit", "rpm_limit"], +) +def test_validate_team_model_max_budget_rejects_unenforceable_entries(model_max_budget, error) -> None: + from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget + + with pytest.raises(HTTPException) as exc: + validate_team_model_max_budget(model_max_budget=model_max_budget, premium_user=True) + assert exc.value.status_code == 400 + assert error in exc.value.detail["error"] + + +def test_validate_team_model_max_budget_accepts_a_zero_cap_and_prefixed_models() -> None: + from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget + + assert ( + validate_team_model_max_budget( + model_max_budget={ + "gpt-4o": BudgetConfig(max_budget=0.0, budget_duration="1d"), + "openai/gpt-4o-mini": BudgetConfig(max_budget=2.5, budget_duration="30d"), + }, + premium_user=True, + ) + is None + ) + + +def test_validate_team_model_max_budget_is_license_gated_only_when_set() -> None: + from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget + + validate_team_model_max_budget(model_max_budget=None, premium_user=False) + validate_team_model_max_budget(model_max_budget={}, premium_user=False) + with pytest.raises(HTTPException) as exc: + validate_team_model_max_budget( + model_max_budget={"gpt-4o": BudgetConfig(max_budget=1.0, budget_duration="1d")}, premium_user=False + ) + assert exc.value.status_code == 403 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 ebbedc6541e..dda5bb344b4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14651,3 +14651,232 @@ async def test_team_info_reports_parent_organization_models_only_to_team_manager ) assert response["team_info"].organization_models == expected_models + + +_EXISTING_TEAM_MODEL_CAPS: Final = { + "gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}, + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "7d"}, +} + + +@pytest.mark.parametrize( + "requested", + [ + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 20.0, "budget_duration": "1d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 10.0, "budget_duration": "30d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"budget_duration": "1d"}}, + {"claude-sonnet-4-6": _EXISTING_TEAM_MODEL_CAPS["claude-sonnet-4-6"]}, + {}, + None, + ], + ids=["raise", "change_duration", "drop_cap_value", "remove_model", "clear_all", "clear_with_null"], +) +def test_team_admin_cannot_loosen_team_model_caps(requested) -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + with pytest.raises(HTTPException) as exc: + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=requested), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"), + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ) + assert exc.value.status_code == 403 + assert "proxy admin" in exc.value.detail["error"] + + +@pytest.mark.parametrize( + "requested", + [ + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 2.0, "budget_duration": "1d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o-mini": {"max_budget": 1.0, "budget_duration": "1d"}}, + dict(_EXISTING_TEAM_MODEL_CAPS), + ], + ids=["lower", "add_model", "unchanged"], +) +def test_team_admin_can_tighten_or_keep_team_model_caps(requested) -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + assert ( + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=requested), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"), + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ) + is None + ) + + +def test_team_model_cap_authority_skips_omitted_field_malformed_rows_and_proxy_admins() -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin") + outcomes = ( + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", max_budget=1.0), + user_api_key_dict=team_admin, + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ), + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget={}), + user_api_key_dict=team_admin, + existing_model_max_budget={"gpt-4o": "not-a-budget", "gpt-4o-mini": {"budget_duration": "1d"}}, + ), + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=None), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ), + ) + assert outcomes == (None, None, None) + + +@pytest.mark.asyncio +async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_auth): + mock_db_client.jsonify_team_object = lambda db_data: db_data + mock_db_client.get_data = AsyncMock(return_value=None) + mock_db_client.update_data = AsyncMock(return_value=MagicMock()) + mock_db_client.db = MagicMock() + mock_db_client.db.litellm_modeltable = MagicMock() + mock_db_client.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) + + team_create_result = MagicMock(team_id="team-model-caps") + team_create_result.model_dump.return_value = {"team_id": "team-model-caps"} + mock_team_create = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + with patch("litellm.proxy.proxy_server.premium_user", True): # test-quality-ok: proxy_server module global is the endpoint's only injection point + await new_team( + data=NewTeamRequest( + team_alias="model-caps", + model_max_budget={"gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}}, + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + team_data = mock_team_create.call_args.kwargs["data"] + assert team_data["model_max_budget"] == { + "gpt-4o": {"max_budget": 10.0, "budget_duration": "1d", "tpm_limit": None, "rpm_limit": None} + } + + +@pytest.mark.asyncio +async def test_new_team_rejects_unenforceable_model_max_budget(mock_db_client, mock_admin_auth): + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException + from litellm.proxy.management_endpoints.team_endpoints import new_team + + mock_db_client.db.litellm_teamtable.create = AsyncMock() + + with patch("litellm.proxy.proxy_server.premium_user", True), pytest.raises(ProxyException) as exc: # test-quality-ok: proxy_server module global is the endpoint's only injection point + await new_team( + data=NewTeamRequest(team_alias="model-caps", model_max_budget={"gpt-4o": {"max_budget": 10.0}}), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + assert exc.value.code == "400" + assert "budget_duration" in str(exc.value.message) + mock_db_client.db.litellm_teamtable.create.assert_not_awaited() + + +def _existing_team_with_model_caps(caps): + existing = MagicMock() + existing.team_id = "standalone-team-123" + existing.organization_id = None + existing.max_budget = None + existing.model_id = None + existing.model_max_budget = caps + existing.model_dump.return_value = { + "team_id": "standalone-team-123", + "organization_id": None, + "model_max_budget": caps, + "members_with_roles": [{"user_id": "team-admin-model-caps", "role": "admin"}], + } + return existing + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cleared_with", [{}, None], ids=["empty_mapping", "null"]) +async def test_update_team_clearing_model_max_budget_writes_an_empty_mapping( + disable_audit_logging_for_mocked_team, cleared_with +): + from fastapi import Request + + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + ): + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=_existing_team_with_model_caps(_EXISTING_TEAM_MODEL_CAPS) + ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + updated = _existing_team_with_model_caps({}) + updated.litellm_model_table = None + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated) + + await update_team( + data=UpdateTeamRequest(team_id="standalone-team-123", model_max_budget=cleared_with), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + assert mock_prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["model_max_budget"] == {} + + +@pytest.mark.asyncio +async def test_update_team_model_max_budget_raise_blocked_for_team_admin(): + from fastapi import Request + + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), # test-quality-ok: stubs the audit write so the test observes only the team update result + ): + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=_existing_team_with_model_caps(_EXISTING_TEAM_MODEL_CAPS) + ) + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.update = AsyncMock() + + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest( + team_id="standalone-team-123", + model_max_budget={ + **_EXISTING_TEAM_MODEL_CAPS, + "gpt-4o": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin-model-caps", models=[] + ), + ) + + assert exc.value.code == "403" + assert "proxy admin" in str(exc.value.message).lower() + mock_prisma.db.litellm_teamtable.update.assert_not_awaited() diff --git a/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py index fe852be775c..0bdf43b396c 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py +++ b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py @@ -47,6 +47,19 @@ def test_team_and_user_state_round_trips_through_metadata(): ) +def test_team_model_max_budget_rides_on_the_token(): + """The team's per-model caps must reach the token, or the auth check and the spend hook never see them.""" + token = UserAPIKeyAuth(token="hashed", team_id="t1") + team_model_max_budget = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", model_max_budget=team_model_max_budget), + user_object=None, + ) + + assert token.team_model_max_budget == team_model_max_budget + + def test_missing_objects_leave_no_metadata_and_no_snapshot(): token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1") carry_team_and_user_budget_state(valid_token=token, team_object=None, user_object=None) diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index ce6ecc2ea65..672dd1eb674 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -401,19 +401,20 @@ async def test_check_view_exists_creates_token_view_when_missing( prisma_client.db.execute_raw = AsyncMock() prisma_client.health_check = AsyncMock(return_value=[{"?column?": 1}]) result = await prisma_client.check_view_exists() + created_sql = prisma_client.db.execute_raw.await_args.args[0] actual = { "result": result, "create_called": prisma_client.db.execute_raw.await_count, - "create_sql_starts_with_create_view": prisma_client.db.execute_raw.await_args.args[ - 0 - ] - .strip() - .startswith('CREATE VIEW "LiteLLM_VerificationTokenView"'), + "create_sql_starts_with_create_view": created_sql.strip().startswith( + 'CREATE VIEW "LiteLLM_VerificationTokenView"' + ), + "projects_team_model_max_budget": "t.model_max_budget AS team_model_max_budget" in created_sql, } assert actual == { "result": None, "create_called": 1, "create_sql_starts_with_create_view": True, + "projects_team_model_max_budget": True, } diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 851c9e6d487..f2d7cff2ec4 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -7,6 +7,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import { toast } from "@/lib/toast"; import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; +import { MODEL_MAX_BUDGET_PREMIUM_HINT } from "./key_team_helpers/ModelMaxBudgetEditor"; import { fetchMCPAccessGroups, getDefaultTeamSettings, @@ -1547,6 +1548,36 @@ describe("Teams - the exact bytes the create call sends", () => { expect(await screen.findByText("Please input a team name")).toBeInTheDocument(); expect(teamCreateCall).not.toHaveBeenCalled(); }); + + it("locks the per-model budget editor and says why when the proxy has no enterprise license", async () => { + await openCreateModal({ premiumUser: false }); + + expect(screen.getByRole("button", { name: /Add Model Budget/i })).toBeDisabled(); + expect(screen.getByText(MODEL_MAX_BUDGET_PREMIUM_HINT)).toBeInTheDocument(); + }); + + it("sends the per-model budget a licensed operator fills in, keyed by model", async () => { + const user = userEvent.setup({ delay: null }); + await openCreateModal({ premiumUser: true }); + + await user.click(screen.getByRole("button", { name: /Add Model Budget/i })); + await chooseSelectOption(user, screen.getByPlaceholderText("Select model"), "gpt-4"); + fireEvent.change(screen.getByPlaceholderText("Max spend ($)"), { target: { value: "3" } }); + + const payload = await submit(); + + expect(payload.model_max_budget).toStrictEqual({ "gpt-4": { budget_limit: 3, time_period: "30d" } }); + }); + + it("leaves model_max_budget out when a started row is removed again", async () => { + const user = userEvent.setup({ delay: null }); + await openCreateModal({ premiumUser: true }); + + await user.click(screen.getByRole("button", { name: /Add Model Budget/i })); + await user.click(screen.getByRole("button", { name: "Remove model budget" })); + + expect(wireBody(await submit())).not.toHaveProperty("model_max_budget"); + }); }); describe("Teams - the create form keeps the organization and models picks while it is open", () => { diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 4f3367d8b98..7214d16f665 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -48,6 +48,7 @@ import BudgetDurationDropdown, { } from "./common_components/budget_duration_dropdown"; import { Organization, getDefaultTeamSettings, getGuardrailsList, getPoliciesList, teamDeleteCall } from "./networking"; import NumericalInput from "./shared/numerical_input"; +import { ModelMaxBudget, ModelMaxBudgetField } from "./key_team_helpers/ModelMaxBudgetEditor"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import SearchToolSelector from "./search_tools/SearchToolSelector"; import SkillSelector from "./skills/SkillSelector"; @@ -271,6 +272,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const [policiesList, setPoliciesList] = useState([]); const [loggingSettings, setLoggingSettings] = useState([]); const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({}); + const [modelMaxBudget, setModelMaxBudget] = useState({}); const [routerSettings, setRouterSettings] = useState(null); const [routerSettingsKey, setRouterSettingsKey] = useState(0); @@ -348,6 +350,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser setSearchToolSettingsOpen(false); setLoggingSettings([]); setModelAliases({}); + setModelMaxBudget({}); setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); }; @@ -525,6 +528,10 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser formValues.model_aliases = modelAliases; } + if (Object.keys(modelMaxBudget).length > 0) { + formValues.model_max_budget = modelMaxBudget; + } + // Add router_settings if any are defined if (routerSettings?.router_settings) { // Only include router_settings if it has at least one non-null value @@ -813,6 +820,14 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser /> )} + {({ ref, value, ...field }) => ( diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx index 0fab5555343..4d2a0e88f84 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx @@ -144,6 +144,7 @@ export function ModelMaxBudgetEditor({ onClick={() => removeEntry(entry.id)} disabled={!premiumUser} title={hintWhenLocked} + aria-label="Remove model budget" className="absolute top-2 right-2 text-muted-foreground hover:text-destructive transition-colors p-1" > diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index eb912ffa3cc..16eb70fdf5e 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1609,6 +1609,99 @@ describe("TeamInfoView", () => { }); }); + describe("per-model budgets", () => { + const teamWithModelBudget = () => + createMockTeamData({ + models: ["gpt-4"], + model_max_budget: { "gpt-4": { max_budget: 5, budget_duration: "1d" } }, + model_max_budget_usage: { "gpt-4": { current_spend: 1.25, budget_limit: 5, time_period: "1d" } }, + }); + + const openSettingsEditor = async (user: ReturnType) => { + await waitFor(() => { + expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0); + }); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await screen.findByLabelText("Team Name"); + }; + + const savedPayload = async () => { + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + return vi.mocked(networking.teamUpdateCall).mock.calls[0][1] as Record; + }; + + it("shows the stored per-model budget and its current spend in the read-only settings view", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0); + }); + await user.click(screen.getByRole("tab", { name: "Settings" })); + + expect(await screen.findByText("Per-Model Budget (gpt-4): $5 per 1d, spent $1.25")).toBeInTheDocument(); + }); + + it("seeds the editor from the stored budget and keeps it read-only without an enterprise license", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + + renderWithProviders(); + + await openSettingsEditor(user); + + expect(screen.getByPlaceholderText("Max spend ($)")).toHaveValue(5); + expect(screen.getByPlaceholderText("Max spend ($)")).toBeDisabled(); + expect(screen.getByRole("button", { name: /Add Model Budget/i })).toBeDisabled(); + }); + + it("leaves model_max_budget out of a save that did not touch it", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect(await savedPayload()).not.toHaveProperty("model_max_budget"); + }); + + it("sends the edited cap for the model", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + fireEvent.change(screen.getByPlaceholderText("Max spend ($)"), { target: { value: "2.5" } }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect((await savedPayload()).model_max_budget).toEqual({ "gpt-4": { budget_limit: 2.5, time_period: "1d" } }); + }); + + it("sends an empty model_max_budget when the last row is removed, so the stored cap is cleared", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + await user.click(screen.getByRole("button", { name: "Remove model budget" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect((await savedPayload()).model_max_budget).toEqual({}); + }); + }); + describe("team member settings", () => { it("should populate Default Key Duration from the team's stored metadata", async () => { const user = userEvent.setup({ delay: null }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index ce705008678..c476f3492a3 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -51,6 +51,13 @@ import GuardrailsSelect from "./GuardrailsSelect"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown"; +import { + ModelBudgetUsage, + ModelMaxBudget, + ModelMaxBudgetField, + modelMaxBudgetToEntries, +} from "../key_team_helpers/ModelMaxBudgetEditor"; +import { modelMaxBudgetUpdate, StoredModelMaxBudget } from "../key_team_helpers/modelMaxBudgetPayload"; import { computeTeamModelBadges, normalizeTeamModelSelection, @@ -268,6 +275,8 @@ export interface TeamData { max_budget: number | null; soft_budget?: number | null; budget_duration: string | null; + model_max_budget?: StoredModelMaxBudget | null; + model_max_budget_usage?: Record | null; models: string[]; blocked: boolean; spend: number; @@ -563,6 +572,7 @@ const TeamInfoView: React.FC = ({ const [isDeleting, setIsDeleting] = useState(false); const [isTeamSaving, setIsTeamSaving] = useState(false); const [teamModelAliases, setTeamModelAliases] = useState>({}); + const [teamModelMaxBudget, setTeamModelMaxBudget] = useState({}); const routerSettingsRef = React.useRef(null); const [organization, setOrganization] = useState(null); const { userRole, userId } = useAuthorized(); @@ -628,6 +638,7 @@ const TeamInfoView: React.FC = ({ const startEditing = () => { form.reset(teamFormValues()); + setTeamModelMaxBudget((teamData?.team_info?.model_max_budget ?? {}) as ModelMaxBudget); setTeamMemberSettingsOpen(false); setSearchToolSettingsOpen(false); setIsEditing(true); @@ -1078,6 +1089,11 @@ const TeamInfoView: React.FC = ({ updateData.model_aliases = teamModelAliases; } + const modelBudgets = modelMaxBudgetUpdate(teamModelMaxBudget, info.model_max_budget); + if (modelBudgets !== undefined) { + updateData.model_max_budget = modelBudgets; + } + // Handle router_settings - read fresh values from DOM at save time. const currentRouterSettings = routerSettingsRef.current?.getValue(); if (currentRouterSettings?.router_settings) { @@ -1536,6 +1552,15 @@ const TeamInfoView: React.FC = ({ )} + + {({ ref, value, ...field }) => } @@ -2051,6 +2076,17 @@ const TeamInfoView: React.FC = ({ : "No Limit"}
Budget Reset: {info.budget_duration || "Never"}
+ {modelMaxBudgetToEntries(info.model_max_budget as ModelMaxBudget | null | undefined).map( + ({ model, budgetLimit, timePeriod }) => { + const spent = model === null ? undefined : info.model_max_budget_usage?.[model]?.current_spend; + return ( +
+ Per-Model Budget ({model}): ${budgetLimit ?? "?"} per {timePeriod} + {spent !== undefined && `, spent $${spent}`} +
+ ); + }, + )} {info.metadata?.soft_budget_alerting_emails && Array.isArray(info.metadata.soft_budget_alerting_emails) && info.metadata.soft_budget_alerting_emails.length > 0 && ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7ca30d5c4f0..f7a26a26370 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -15626,6 +15626,7 @@ export interface paths { * - prompts: Optional[List[str]] - List of prompts that the team is allowed to use. * - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. * - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + * - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} * - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) * - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -15852,6 +15853,7 @@ export interface paths { * - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). * - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. * - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + * - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} * - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) * - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -33328,6 +33330,13 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** + * Model Max Budget + * @description Max budget per model for every key on the team, overridable per key (e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}}) + */ + model_max_budget?: { + [key: string]: components["schemas"]["BudgetConfig"]; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -34086,6 +34095,13 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** + * Model Max Budget + * @description Max budget per model for every key on the team, overridable per key (e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}}) + */ + model_max_budget?: { + [key: string]: components["schemas"]["BudgetConfig"]; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -39225,6 +39241,13 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** + * Model Max Budget + * @description Max budget per model for every key on the team, overridable per key (e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}}) + */ + model_max_budget?: { + [key: string]: components["schemas"]["BudgetConfig"]; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -39936,6 +39959,10 @@ export interface components { team_model_aliases?: { [key: string]: unknown; } | null; + /** Team Model Max Budget */ + team_model_max_budget?: { + [key: string]: unknown; + } | null; /** * Team Models * @default [] From b3432abef7920773788382709bc160c4c0c7b9e9 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 00:58:03 +0000 Subject: [PATCH 042/187] refactor(ui): derive the ssh skill name from the bare repo path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../components/claude_code_plugins/helpers.ts | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts index cbca22da1f3..3c11a6cedbc 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts @@ -180,17 +180,15 @@ const parseRawGitSource = (url: URL, subPath?: string): SkillSourcePreview | nul interface SshRemote { cloneUrl: string; - repoPath: string; + repoName: string; } -const withGitSuffix = (path: string): string => `${path.replace(/\.git$/i, "")}.git`; - const buildSshRemote = (rawPath: string, toCloneUrl: (repoPath: string) => string): SshRemote | null => { if (rawPath.split("/").some((segment) => DOTS_ONLY_SEGMENT_REGEX.test(segment))) { return null; } - const repoPath = withGitSuffix(rawPath); - return { cloneUrl: toCloneUrl(repoPath), repoPath }; + const bare = rawPath.replace(/\.git$/i, ""); + return { cloneUrl: toCloneUrl(`${bare}.git`), repoName: lastSegment(bare) }; }; const parseSshRemote = (raw: string): SshRemote | null => { @@ -208,9 +206,6 @@ const parseSshRemote = (raw: string): SshRemote | null => { return null; }; -const parseSshSource = (remote: SshRemote, subPath?: string): SkillSourcePreview | null => - buildGitSourcePreview("SSH", remote.cloneUrl, lastSegment(remote.repoPath).replace(/\.git$/, ""), subPath); - const parseArchiveSource = (url: URL): SkillSourcePreview => ({ parsed: { source: "archive", url: url.href }, label: `Zip archive — ${url.host}${url.pathname}`, @@ -225,9 +220,9 @@ const parseArchiveSource = (url: URL): SkillSourcePreview => ({ * with an optional subfolder turning it into git-subdir. */ export const parseSkillSource = (rawUrl: string, subPath?: string): SkillSourcePreview | null => { - const sshRemote = parseSshRemote(rawUrl); - if (sshRemote) { - return parseSshSource(sshRemote, subPath); + const ssh = parseSshRemote(rawUrl); + if (ssh) { + return buildGitSourcePreview("SSH", ssh.cloneUrl, ssh.repoName, subPath); } const url = parseRepoUrl(rawUrl); if (!url) { From 3e8566d87849b20299f2a07889e2dfe5ab35cab6 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 01:03:40 +0000 Subject: [PATCH 043/187] fix(keys): keep organization_id on archived key records and /key/info Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/models/verification_token.py | 1 + .../test_key_management_endpoints.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index 06ff877a41a..1b807c46c40 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -69,6 +69,7 @@ class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): """Audit record for deleted keys; mirrors the token plus deletion metadata.""" id: str | None = None + organization_id: str | None = None deleted_at: datetime | None = None deleted_by: str | None = None deleted_by_api_key: str | None = None diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index a8ff860c7c9..34cfc2a8fac 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -4919,6 +4919,23 @@ def test_transform_verification_tokens_to_deleted_records(): assert json.loads(record2["budget_fallbacks"]) == {"gpt-4": ["gpt-4o-mini"]} +def test_transform_verification_tokens_to_deleted_records_keeps_organization_id(): + live_row = MagicMock() + live_row.model_dump.return_value = { + "token": "hashed-token-org", + "user_id": "user-123", + "team_id": None, + "organization_id": "org-finops", + } + + records = _transform_verification_tokens_to_deleted_records( + keys=[live_row], + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", api_key="sk-admin"), + ) + + assert records[0]["organization_id"] == "org-finops" + + def test_transform_verification_tokens_to_deleted_records_empty_list(): user_api_key_dict = UserAPIKeyAuth( user_id="user-123", @@ -6135,6 +6152,7 @@ def _archived_key_row(token: str, user_id: str) -> MagicMock: "key_alias": "finops-2024", "user_id": user_id, "team_id": None, + "organization_id": "org-finops", "blocked": None, "deleted_at": datetime(2024, 11, 15, 10, 0, tzinfo=timezone.utc), "deleted_by": "admin-1", @@ -6166,6 +6184,7 @@ async def test_info_key_fn_serves_deleted_key_from_archive(monkeypatch): info = result["info"] assert info["status"] == "deleted" assert info["key_alias"] == "finops-2024" + assert info["organization_id"] == "org-finops" assert info["deleted_by"] == "admin-1" assert info["deleted_at"] is not None assert "token" not in info From 13553473aaf1afaac6157ff8d0c87a1984660c8e Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:06:00 -0700 Subject: [PATCH 044/187] fix(mcp): reject missing upstream authentication credentials --- litellm/experimental_mcp_client/client.py | 8 +- .../mcp_server/mcp_server_manager.py | 87 ++++---- .../mcp_server/openapi_to_mcp_generator.py | 24 +-- .../outbound_credentials/adapter.py | 24 +-- .../proxy/_experimental/mcp_server/server.py | 1 + .../_experimental/mcp_server/upstream.py | 81 ++++++++ .../proxy/_experimental/mcp_server/utils.py | 16 ++ .../test_mcp_client.py | 15 ++ .../outbound_credentials/test_adapter.py | 17 +- .../mcp_server/test_mcp_hook_extra_headers.py | 1 + .../mcp_server/test_mcp_server_manager.py | 190 +++++++++++++++++- .../test_openapi_to_mcp_generator.py | 18 ++ 12 files changed, 396 insertions(+), 86 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/upstream.py diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index ee01a53ecb3..56ee5f30d02 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -346,13 +346,17 @@ class MCPClient: self.update_auth_value(auth_value) async def discovery_auth_fingerprint(self) -> str: + return self._hash_discovery_auth(await self.prepare_request_auth()) + + async def prepare_request_auth(self) -> httpx.Request: + """Preview the authenticated request without sending it, closing the auth flow afterwards.""" request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers()) if self._resolved_auth is None: - return self._hash_discovery_auth(request) + return request flow: Final = self._resolved_auth.async_auth_flow(request) try: authenticated: Final = await flow.__anext__() - return self._hash_discovery_auth(authenticated) + return authenticated finally: await flow.aclose() diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index fb0c623473a..0254f79cbcc 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -132,6 +132,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) +from litellm.proxy._experimental.mcp_server.upstream import prepare_mcp_client, validate_openapi_credentials from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, MCPMissingUserEnvVarsError, @@ -4229,16 +4230,19 @@ class MCPServerManager: ) record_auth_resolution(server.server_id, AuthResolution.not_applicable) - return MCPClient( - server_url="", # Not used for stdio - transport_type=transport, - auth_type=resolved_server.auth_type, - auth_value=auth_value, - timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), - stdio_config=stdio_config, - extra_headers=extra_headers, - sampling_callback=sampling_cb, - elicitation_callback=elicitation_cb, + return await prepare_mcp_client( + resolved_server, + MCPClient( + server_url="", # Not used for stdio + transport_type=transport, + auth_type=resolved_server.auth_type, + auth_value=auth_value, + timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), + stdio_config=stdio_config, + extra_headers=extra_headers, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, + ), ) else: # For HTTP/SSE transports @@ -4259,15 +4263,20 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, extra_headers=extra_headers, ) - return MCPClient( - server_url=server_url, - transport_type=transport, - auth_type=resolved_server.auth_type, - timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), - extra_headers=extra_headers, - resolved_auth=resolved_auth, - sampling_callback=sampling_cb, - elicitation_callback=elicitation_cb, + return await prepare_mcp_client( + resolved_server, + MCPClient( + server_url=server_url, + transport_type=transport, + auth_type=resolved_server.auth_type, + timeout=( + resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT + ), + extra_headers=extra_headers, + resolved_auth=resolved_auth, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, + ), ) # Create SigV4 auth if configured @@ -4297,17 +4306,20 @@ class MCPServerManager: else AuthResolution.no_auth ) record_auth_resolution(server.server_id, legacy_source) - return MCPClient( - server_url=server_url, - transport_type=transport, - auth_type=resolved_server.auth_type, - auth_value=auth_value, - auth_header_name=auth_header_name, - timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), - extra_headers=extra_headers, - aws_auth=aws_auth, - sampling_callback=sampling_cb, - elicitation_callback=elicitation_cb, + return await prepare_mcp_client( + resolved_server, + MCPClient( + server_url=server_url, + transport_type=transport, + auth_type=resolved_server.auth_type, + auth_value=auth_value, + auth_header_name=auth_header_name, + timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), + extra_headers=extra_headers, + aws_auth=aws_auth, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, + ), ) async def _get_tools_from_server( @@ -6188,6 +6200,7 @@ class MCPServerManager: mcp_auth_header: str | dict[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None, forwarded_headers: dict[str, str] | None, + caller_authorization: str | None = None, ) -> tuple[dict[str, str] | None, dict[str, str] | None]: """Resolve the gateway-owned upstream credential for a spec_path (OpenAPI) tool call. @@ -6211,9 +6224,12 @@ class MCPServerManager: """ spec: Final = to_server_spec(mcp_server) if spec is None: - if oauth2_headers: - return None, forwarded_headers - stored_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) + stored_headers = ( + None + if oauth2_headers + else await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) + ) + validate_openapi_credentials(mcp_server, stored_headers, forwarded_headers, caller_authorization) return stored_headers, forwarded_headers subject_token: str | None = None @@ -6232,7 +6248,9 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, extra_headers=forwarded_headers, ) - return await _materialize_auth_headers(resolved_auth), forwarded_headers + resolved_headers: Final = await _materialize_auth_headers(resolved_auth) + validate_openapi_credentials(mcp_server, resolved_headers, forwarded_headers, caller_authorization) + return resolved_headers, forwarded_headers async def _gather_openapi_tool_tasks( self, @@ -6358,6 +6376,7 @@ class MCPServerManager: mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, + caller_authorization=auth_header_value, ) async def _call_openapi_via_handler(): diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index d115eb8b3c1..66712e97a34 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -20,6 +20,7 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( MCPOpenApiUpstreamError, MCPUpstreamAuthError, ) +from litellm.proxy._experimental.mcp_server.utils import merge_openapi_headers # Tool names emitted from OpenAPI specs must work across all major LLM providers. # OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to @@ -415,26 +416,9 @@ def _merge_openapi_tool_request_headers( Header names are compared case-insensitively so different casing cannot bypass the precedence rules. """ - request_extra: Final = _request_extra_headers.get() or {} - static: Final = static_headers or {} - - static_lower_names: Final = {k.lower() for k in static} - effective_headers: dict[str, str] = {k: v for k, v in request_extra.items() if k.lower() not in static_lower_names} - effective_headers.update(static) - - override_auth: Final = _request_auth_header.get() - if override_auth: - for existing in [k for k in effective_headers if k.lower() == "authorization"]: - del effective_headers[existing] - effective_headers["Authorization"] = override_auth - - resolved_auth_headers: Final = _request_resolved_auth_headers.get() or {} - for name, value in resolved_auth_headers.items(): - for existing in [k for k in effective_headers if k.lower() == name.lower()]: - del effective_headers[existing] - effective_headers[name] = value - - return effective_headers + return merge_openapi_headers( + static_headers, _request_extra_headers.get(), _request_auth_header.get(), _request_resolved_auth_headers.get() + ) def _raise_for_upstream_failure( diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 77979a15199..d25946d81d0 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -79,7 +79,7 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None: BYOK is the per-user source of the ``api_key`` mode; its scheme rides on ``auth_type`` just like a shared key, but the value is per-user and not migrated yet, so a BYOK server defers - to v1 regardless of ``auth_type`` (this guard is the seam the BYOK arm replaces later). + to v1 for its static schemes. Declared OBO always stays with the exchange arm. Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is @@ -90,8 +90,8 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None: modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); delegated/passthrough oauth2 and SigV4 return None and stay on v1. """ - if server.is_byok: - return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type) + if server.is_byok and server.auth_type != MCPAuth.oauth2_token_exchange: + return None # per-user BYOK source not migrated yet -> defer to v1 resource: Final = server.url or server.server_id auth_type: Final = server.auth_type match auth_type: @@ -165,21 +165,9 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec: ) -def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None: - """Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured. - - An OBO server with ``client_id``/``client_secret`` is owned by the v2 arm even if the - ``token_exchange_endpoint``/``token_url`` is absent: a missing endpoint then fails closed (412) at - the exchanger rather than silently deferring to v1 and connecting unauthenticated, since the - gateway must not guess the IdP or fall back to a weaker source. Without client credentials there is - nothing to own, so the server stays on v1 (parity-safe). ``profile`` selects the wire dialect - (``rfc8693`` default, ``entra_obo`` for Microsoft Entra On-Behalf-Of); an unrecognized value - normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is - forwarded only when the operator set it; a missing one is omitted, not derived. - """ +def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec: + """Keep declared OBO owned by the resolver, including incomplete client configuration.""" endpoint: Final = server.token_exchange_endpoint or server.effective_token_url - if not server.client_id or not server.client_secret: - return None profile: Final[Literal["rfc8693", "entra_obo"]] = ( "entra_obo" if server.token_exchange_profile == "entra_obo" else "rfc8693" ) @@ -193,7 +181,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None: token_exchange_endpoint=endpoint, audience=server.audience, client_id=server.client_id, - client_secret=SecretStr(server.client_secret), + client_secret=SecretStr(server.client_secret) if server.client_secret else None, token_endpoint_auth_method=server.token_endpoint_auth_method, scopes=tuple(server.scopes or ()), ), diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7feb1fd468d..a3aaada41f7 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3141,6 +3141,7 @@ if MCP_AVAILABLE: mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, + caller_authorization=auth_header_value, ) _auth_token: Final = _request_auth_header.set(auth_header_value) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py new file mode 100644 index 00000000000..8d7a7dba252 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from typing import Final + +from litellm.experimental_mcp_client.client import MCPClient +from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import raise_public +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result +from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError +from litellm.proxy._experimental.mcp_server.utils import merge_openapi_headers +from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +_STATIC_MODES: Final = frozenset( + (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization) +) + + +def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool: + if not value: + return False + if auth_type == MCPAuth.authorization or (auth_type == MCPAuth.api_key and name != "authorization"): + return True + if value.lower() in ("bearer", "basic", "token", "apikey"): + return False + if auth_type == MCPAuth.basic: + parts: Final = value.split(None, 1) + if len(parts) != 2 or parts[0].lower() != "basic": + return False + try: + return bool(base64.b64decode(parts[1], validate=True).strip()) + except ValueError: + return False + return True + + +def validate_static_credential( + server: MCPServer, headers: Mapping[str, str], *, header_slot: str | None = None, openapi: bool = False +) -> Result[None, CredError]: + if server.auth_type not in _STATIC_MODES or server.transport == MCPTransport.stdio: + return Ok(None) + default_slot: Final = "X-API-Key" if server.auth_type == MCPAuth.api_key else "Authorization" + slots: Final = frozenset( + name.lower() + for name in ( + header_slot or server.upstream_token_header or default_slot, + "Authorization" if openapi else default_slot, + ) + ) + values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) + if values and all(_usable_credential_value(server.auth_type, name, value) for name, value in values): + return Ok(None) + return Error(CredError.of_misconfigured(f"{server.auth_type} requires a usable upstream credential")) + + +async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient: + if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio: + return client + request: Final = await client.prepare_request_auth() + match validate_static_credential(server, request.headers): + case Error(error): + raise_public(error) + case Ok(): + return client + + +def validate_openapi_credentials( + server: MCPServer, + resolved_headers: Mapping[str, str] | None, + forwarded_headers: Mapping[str, str] | None, + caller_authorization: str | None, +) -> None: + headers: Final = merge_openapi_headers( + server.static_headers or {}, forwarded_headers, caller_authorization, resolved_headers + ) + match validate_static_credential(server, headers, openapi=True): + case Error(error): + raise_public(error) + case Ok(): + return diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index fb3eb06fd15..bea74d36b34 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -756,6 +756,22 @@ def build_env_var_setup_url(server_id: str) -> str: return f"{base}{path}" if base else path +def merge_openapi_headers( + static_headers: Mapping[str, str], + extra_headers: Mapping[str, str] | None, + caller_authorization: str | None, + resolved_headers: Mapping[str, str] | None, +) -> dict[str, str]: + sources: Final = ( + extra_headers or {}, + static_headers, + {"Authorization": caller_authorization} if caller_authorization else {}, + resolved_headers or {}, + ) + entries: Final = {name.lower(): (name, value) for source in sources for name, value in source.items()} + return dict(entries.values()) + + def merge_mcp_headers( *, extra_headers: Mapping[str, str] | None = None, diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index f72316f5d5e..d9ffb0d64fe 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1934,3 +1934,18 @@ async def test_discovery_auth_fingerprint_tracks_effective_credentials(resolved: assert original != replaced assert len(original) == 64 assert "private-original-credential" not in original + + +@pytest.mark.asyncio +async def test_request_auth_preview_uses_the_same_effective_headers_as_egress() -> None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth + + client: Final = MCPClient( + server_url="https://upstream.example/mcp", auth_type=MCPAuth.bearer_token, + resolved_auth=StaticHeaderAuth("Bearer resolved"), extra_headers={"X-Trace": "trace"}, + ) + request: Final = await client.prepare_request_auth() + assert request.method == "POST" + assert str(request.url) == "https://upstream.example/mcp" + assert request.headers["Authorization"] == "Bearer resolved" + assert request.headers["X-Trace"] == "trace" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 1b003e11993..2885fdaef95 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -155,12 +155,6 @@ def test_oauth2_user_token_maps_to_authorization_code(oauth2_flow): _server(auth_type=MCPAuth.api_key), # no token configured _server(auth_type=MCPAuth.bearer_token), # no token configured _server(auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True), # delegated upstream OAuth -> v1 - _server(auth_type=MCPAuth.oauth2_token_exchange), # no endpoint/client creds -> incomplete -> v1 - _server( - auth_type=MCPAuth.oauth2_token_exchange, - token_exchange_endpoint="https://idp/token", - client_id="cid", - ), # missing client_secret -> incomplete -> v1 _server(auth_type=MCPAuth.aws_sigv4), _server(auth_type=None, oauth_passthrough=True, extra_headers=["Authorization"]), ], @@ -802,3 +796,14 @@ def test_a_blank_header_name_means_unset_rather_than_an_error(blank): spec = to_server_spec(server) assert spec is not None assert spec.config.header_name == "Authorization" + + +@pytest.mark.parametrize("client_secret", [None, ""]) +@pytest.mark.parametrize("is_byok", [False, True]) +def test_incomplete_obo_keeps_exchange_ownership(client_secret: str | None, is_byok: bool) -> None: + spec = to_server_spec(_server(auth_type=MCPAuth.oauth2_token_exchange, client_id="client", + client_secret=client_secret, is_byok=is_byok)) + assert spec is not None + assert isinstance(spec.config, TokenExchangeConfig) + assert spec.config.client_id == "client" + assert spec.config.client_secret is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 28faf375ab8..5e3a26fb4ac 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1391,6 +1391,7 @@ class TestOpenApiResolvedUpstreamAuth: mcp_auth_header="user-byok-key", user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="sk-user"), forwarded_headers=None, + caller_authorization="ApiKey user-byok-key", ) assert resolved is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d56f08c4e79..ea05035e16f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -9401,12 +9401,13 @@ class TestCreateMcpClientV2Graft: assert "misconfigured" in str(exc_info.value.detail) assert "token_url" in str(exc_info.value.detail) - async def test_static_token_missing_defers_to_v1(self): - client = await MCPServerManager()._create_mcp_client( - self._http_server(auth_type=MCPAuth.api_key, authentication_token=None) - ) - - assert client._resolved_auth is None + async def test_static_token_missing_rejects_before_connecting(self): + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client( + self._http_server(auth_type=MCPAuth.api_key, authentication_token=None) + ) + assert exc.value.status_code == 500 + assert "credential" in str(exc.value.detail) async def test_stdio_migrated_auth_type_still_defers_to_v1(self): client = await MCPServerManager()._create_mcp_client( @@ -13467,3 +13468,180 @@ async def test_discovery_cache_returns_oversized_results_without_retaining_them( result: Final = await cache.get(("server", None), fetch) assert result[0].description == description assert fetch.await_count == 2 + + +class TestProtectedCredentialPreparation: + @pytest.mark.asyncio + @pytest.mark.parametrize("transport", [MCPTransport.http, MCPTransport.sse]) + @pytest.mark.parametrize("client_secret", [None, ""]) + @pytest.mark.parametrize("subject", [None, "caller-subject"]) + async def test_incomplete_obo_rejects_caller_and_static_fallback( + self, transport: MCPTransport, client_secret: str | None, subject: str | None + ) -> None: + server = MCPServer( + server_id="incomplete-obo", name="incomplete-obo", url="https://upstream.example/mcp", + transport=transport, auth_type=MCPAuth.oauth2_token_exchange, + client_id="gateway", client_secret=client_secret, + token_exchange_endpoint="https://idp.example/token", authentication_token="static-fallback", + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client( + server, mcp_auth_header="Bearer override", subject_token=subject, + ) + assert exc.value.status_code == (401 if subject is None else 500) + assert "static-fallback" not in str(exc.value.detail) + assert "override" not in str(exc.value.detail) + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.api_key, MCPAuth.bearer_token]) + @pytest.mark.parametrize("credential", [None, "", " ", {"X-Trace": "trace"}]) + async def test_static_auth_without_usable_credential_rejects( + self, auth_type: MCPAuthType, credential: str | dict[str, str] | None + ) -> None: + server = MCPServer( + server_id="empty-static", name="empty-static", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header=credential) + assert exc.value.status_code == 500 + assert "credential" in str(exc.value.detail).lower() + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,headers", [ + (MCPAuth.api_key, {"X-API-Key": "key"}), + (MCPAuth.bearer_token, {"Authorization": "Bearer token"}), + ]) + async def test_static_auth_accepts_actual_forwarded_credential( + self, auth_type: MCPAuthType, headers: dict[str, str] + ) -> None: + server = MCPServer( + server_id="header-static", name="header-static", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, + ) + client = await MCPServerManager()._create_mcp_client(server, extra_headers=headers) + assert client._get_auth_headers() == headers + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.api_key, MCPAuth.bearer_token]) + async def test_openapi_protected_auth_rejects_missing_credentials(self, auth_type: MCPAuthType) -> None: + server = MCPServer( + server_id="openapi-empty", name="openapi-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, + token_exchange_endpoint="https://idp.example/token", + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager().resolve_openapi_upstream_auth( + mcp_server=server, oauth2_headers=None, raw_headers=None, mcp_auth_header=None, + user_api_key_auth=None, forwarded_headers=None, + ) + assert exc.value.status_code in (401, 500) + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,slot", [(MCPAuth.api_key, "X-API-Key"), (MCPAuth.authorization, "Authorization")]) + async def test_raw_static_value_named_token_is_a_usable_credential(self, auth_type: MCPAuthType, slot: str) -> None: + server = MCPServer(server_id="raw-key", name="raw-key", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, authentication_token="token") + client = await MCPServerManager()._create_mcp_client(server) + assert client._resolved_auth is not None + request = httpx.Request("GET", server.url) + flow = client._resolved_auth.auth_flow(request) + try: + assert next(flow).headers[slot] == "token" + finally: + flow.close() + + @pytest.mark.asyncio + async def test_byok_flag_cannot_bypass_incomplete_obo(self) -> None: + server = MCPServer(server_id="obo-byok", name="obo-byok", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.oauth2_token_exchange, is_byok=True, + token_exchange_endpoint="https://idp.example/token") + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header="Bearer override") + assert exc.value.status_code == 401 + + @pytest.mark.asyncio + @pytest.mark.parametrize("configured,override", [(None, "Bearer usable"), ("shared", "Bearer usable")]) + async def test_bearer_override_remains_usable(self, configured: str | None, override: str) -> None: + server = MCPServer(server_id="override", name="override", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=configured) + client = await MCPServerManager()._create_mcp_client(server, mcp_auth_header=override) + assert client._get_auth_headers()["Authorization"] == override + + @pytest.mark.asyncio + @pytest.mark.parametrize("token", [None, "shared"]) + async def test_empty_injected_header_cannot_satisfy_protected_auth(self, token: str | None) -> None: + server = MCPServer(server_id="empty-header", name="empty-header", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=token) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, extra_headers={"authorization": " "}) + assert exc.value.status_code == 500 + + @pytest.mark.asyncio + async def test_custom_slot_uses_its_actual_credential(self) -> None: + server = MCPServer(server_id="custom", name="custom", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, + upstream_token_header="X-Custom", authentication_token="key") + client = await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Trace": "trace"}) + assert client._credential_slot == "X-Custom" + assert await client.discovery_auth_fingerprint() + + @pytest.mark.asyncio + @pytest.mark.parametrize("static,forwarded,caller", [ + ({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None), + ({}, {"X-API-Key": "forwarded"}, None), + ({}, None, "ApiKey caller"), + ]) + async def test_openapi_static_credentials_remain_supported( + self, static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None + ) -> None: + server = MCPServer(server_id="openapi-static", name="openapi-static", url="https://upstream.example", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, static_headers=static) + resolved, retained = await MCPServerManager().resolve_openapi_upstream_auth( + mcp_server=server, oauth2_headers=None, raw_headers=None, mcp_auth_header=None, + user_api_key_auth=None, forwarded_headers=forwarded, caller_authorization=caller, + ) + assert resolved is None + assert retained == forwarded + + @pytest.mark.asyncio + async def test_static_resolution_cancellation_closes_flow(self) -> None: + from collections.abc import AsyncGenerator + from litellm.experimental_mcp_client.client import MCPClient + from litellm.proxy._experimental.mcp_server.upstream import prepare_mcp_client + + class CancelledAuth(httpx.Auth): + closed = False + + async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]: + try: + raise asyncio.CancelledError() + yield request + finally: + self.closed = True + + auth = CancelledAuth() + server = MCPServer(server_id="cancel", name="cancel", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key) + client = MCPClient(server_url=server.url, auth_type=MCPAuth.api_key, resolved_auth=auth) + with pytest.raises(asyncio.CancelledError): + await prepare_mcp_client(server, client) + assert auth.closed + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type", [MCPAuth.basic, MCPAuth.token, MCPAuth.authorization]) + async def test_other_static_schemes_reject_whitespace_credentials(self, auth_type: MCPAuthType) -> None: + server = MCPServer(server_id="blank-static", name="blank-static", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=" ") + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server) + assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc"]) + async def test_basic_headers_without_usable_credentials_reject(self, header: str) -> None: + server = MCPServer(server_id="bad-basic", name="bad-basic", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, extra_headers={"Authorization": header}) + assert exc.value.status_code == 500 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 5fa202224e3..66c5627bc94 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -1458,3 +1458,21 @@ class TestBoundedOpenAPISpecLoading: else: assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) == {"paths": {}} assert destination.call_count == 1 + + +def test_openapi_generator_import_does_not_require_mcp_sdk() -> None: + import subprocess + import sys + + script = """ +import builtins +original_import = builtins.__import__ +def without_mcp(name, *args, **kwargs): + if name == 'mcp' or name.startswith('mcp.'): + raise ModuleNotFoundError('MCP SDK unavailable') + return original_import(name, *args, **kwargs) +builtins.__import__ = without_mcp +import litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator +""" + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True) + assert result.returncode == 0, result.stderr From d8ef940232b906001113d6aaf35ded908213437a Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 01:06:36 +0000 Subject: [PATCH 045/187] chore(ui): regenerate schema.d.ts for organization_id on archived key records Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 542d491e2d7..49c67dcafc8 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29528,6 +29528,8 @@ export interface components { object_permission_id?: string | null; /** Org Id */ org_id?: string | null; + /** Organization Id */ + organization_id?: string | null; /** * Permissions * @default {} From d90e7b3aecb2f9494df669de90f80037e80bc8c5 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 01:26:40 +0000 Subject: [PATCH 046/187] fix(team): resolve model aliases in team admin model_max_budget authority check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 43 +++++++++++-------- .../test_team_endpoints.py | 18 +++++++- 2 files changed, 42 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6fb1ef5ec93..23b02c22f24 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -96,7 +96,10 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage +from litellm.proxy.hooks.model_max_budget_limiter import ( + build_model_max_budget_usage, + resolve_model_budget, +) from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity_aggregated, ) @@ -1194,31 +1197,37 @@ def _check_team_model_budget_update_authority( requested: Final[Mapping[str, BudgetConfig]] = data.model_max_budget or {} for model_name, raw_existing in existing_model_max_budget.items(): existing = _existing_model_cap(raw_existing) - if existing is None or existing.max_budget is None: + if existing is None or existing.max_budget is None or model_name in requested: + continue + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"Only a proxy admin can remove a team's model_max_budget for {model_name!r}. " + f"Current max_budget={existing.max_budget}." + ) + }, + ) + for model_name, proposed in requested.items(): + governing = resolve_model_budget(model=model_name, model_max_budget=existing_model_max_budget) + if governing is None: + continue + cap = governing.budget_config + if cap.max_budget is None: continue - proposed = requested.get(model_name) - if proposed is None: - raise HTTPException( - status_code=403, - detail={ - "error": ( - f"Only a proxy admin can remove a team's model_max_budget for {model_name!r}. " - f"Current max_budget={existing.max_budget}." - ) - }, - ) if ( proposed.max_budget is None - or proposed.max_budget > existing.max_budget - or proposed.budget_duration != existing.budget_duration + or proposed.max_budget > cap.max_budget + or proposed.budget_duration != cap.budget_duration ): raise HTTPException( status_code=403, detail={ "error": ( f"Only a proxy admin can raise a team's model_max_budget for {model_name!r} or change its " - f"budget_duration. Current max_budget={existing.max_budget} per {existing.budget_duration}, " - f"requested={proposed.max_budget} per {proposed.budget_duration}." + f"budget_duration. Current max_budget={cap.max_budget} per {cap.budget_duration} " + f"(entry {governing.budget_model!r}), requested={proposed.max_budget} per " + f"{proposed.budget_duration}." ) }, ) 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 dda5bb344b4..6478b18e553 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14668,8 +14668,21 @@ _EXISTING_TEAM_MODEL_CAPS: Final = { {"claude-sonnet-4-6": _EXISTING_TEAM_MODEL_CAPS["claude-sonnet-4-6"]}, {}, None, + {**_EXISTING_TEAM_MODEL_CAPS, "openai/gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "openai/gpt-4o": {"max_budget": 10.0, "budget_duration": "30d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "anthropic/claude-sonnet-4-6": {"budget_duration": "7d"}}, + ], + ids=[ + "raise", + "change_duration", + "drop_cap_value", + "remove_model", + "clear_all", + "clear_with_null", + "raise_via_provider_alias", + "rewindow_via_provider_alias", + "uncap_via_provider_alias", ], - ids=["raise", "change_duration", "drop_cap_value", "remove_model", "clear_all", "clear_with_null"], ) def test_team_admin_cannot_loosen_team_model_caps(requested) -> None: from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority @@ -14690,8 +14703,9 @@ def test_team_admin_cannot_loosen_team_model_caps(requested) -> None: {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 2.0, "budget_duration": "1d"}}, {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o-mini": {"max_budget": 1.0, "budget_duration": "1d"}}, dict(_EXISTING_TEAM_MODEL_CAPS), + {**_EXISTING_TEAM_MODEL_CAPS, "openai/gpt-4o": {"max_budget": 2.0, "budget_duration": "1d"}}, ], - ids=["lower", "add_model", "unchanged"], + ids=["lower", "add_model", "unchanged", "tighten_via_provider_alias"], ) def test_team_admin_can_tighten_or_keep_team_model_caps(requested) -> None: from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority From 8a3add3c6a7ad900e19d2ed7760627612aa3d17d Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:27:25 -0700 Subject: [PATCH 047/187] fix(mcp): reject scheme-only Basic credentials --- litellm/proxy/_experimental/mcp_server/upstream.py | 3 ++- .../mcp_server/test_mcp_server_manager.py | 13 ++++++++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index 8d7a7dba252..89fd1ad5066 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -29,7 +29,8 @@ def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> b if len(parts) != 2 or parts[0].lower() != "basic": return False try: - return bool(base64.b64decode(parts[1], validate=True).strip()) + decoded: Final = base64.b64decode(parts[1], validate=True).strip() + return bool(decoded) and decoded.lower() != b"basic" except ValueError: return False return True diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ea05035e16f..e2275d05fcf 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13638,10 +13638,21 @@ class TestProtectedCredentialPreparation: assert exc.value.status_code == 500 @pytest.mark.asyncio - @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc"]) + @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc", "Basic QmFzaWM="]) async def test_basic_headers_without_usable_credentials_reject(self, header: str) -> None: server = MCPServer(server_id="bad-basic", name="bad-basic", url="https://upstream.example/mcp", transport=MCPTransport.http, auth_type=MCPAuth.basic) with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, extra_headers={"Authorization": header}) assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["Basic", "Basic ", "basic"]) + @pytest.mark.parametrize("source", ["configured", "caller"]) + async def test_basic_scheme_alone_is_not_a_credential(self, value: str, source: str) -> None: + server = MCPServer(server_id="basic-scheme", name="basic-scheme", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic, + authentication_token=value if source == "configured" else None) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None) + assert exc.value.status_code == 500 From 84e14789d1c7861fbfe76534c6eb091bcaf00f7b Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:18:48 -0700 Subject: [PATCH 048/187] fix(mcp): preserve usable alternate header credentials --- .../_experimental/mcp_server/upstream.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 37 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index 89fd1ad5066..49aca5cf31b 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -50,7 +50,7 @@ def validate_static_credential( ) ) values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) - if values and all(_usable_credential_value(server.auth_type, name, value) for name, value in values): + if any(_usable_credential_value(server.auth_type, name, value) for name, value in values): return Ok(None) return Error(CredError.of_misconfigured(f"{server.auth_type} requires a usable upstream credential")) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e2275d05fcf..c8dba1c1554 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13591,6 +13591,7 @@ class TestProtectedCredentialPreparation: ({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None), ({}, {"X-API-Key": "forwarded"}, None), ({}, None, "ApiKey caller"), + ({"X-API-Key": "static"}, {"Authorization": ""}, None), ]) async def test_openapi_static_credentials_remain_supported( self, static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None @@ -13656,3 +13657,39 @@ class TestProtectedCredentialPreparation: with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None) assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,value,default_slot", [ + (MCPAuth.api_key, "fixture-key", "X-API-Key"), + (MCPAuth.bearer_token, "fixture-key", "Authorization"), + (MCPAuth.basic, "user:pass", "Authorization"), + (MCPAuth.token, "fixture-key", "Authorization"), + (MCPAuth.authorization, "fixture-key", "Authorization"), + ]) + @pytest.mark.parametrize("source", ["configured", "caller"]) + async def test_usable_credential_survives_an_empty_alternate_header( + self, auth_type: MCPAuthType, value: str, default_slot: str, source: str + ) -> None: + server: Final = MCPServer( + server_id="alternate", name="alternate", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, upstream_token_header="X-Custom", + authentication_token=value if source == "configured" else None, + ) + empty_slot: Final = default_slot if source == "configured" else "X-Custom" + selected_slot: Final = "X-Custom" if source == "configured" else default_slot + client: Final = await MCPServerManager()._create_mcp_client( + server, mcp_auth_header=value if source == "caller" else None, extra_headers={empty_slot: ""}, + ) + request: Final = await client.prepare_request_auth() + assert request.headers[selected_slot] + assert request.headers[empty_slot] == "" + + @pytest.mark.asyncio + async def test_empty_custom_and_default_headers_do_not_satisfy_auth(self) -> None: + server: Final = MCPServer( + server_id="both-empty", name="both-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, upstream_token_header="X-Custom", + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Custom": "", "X-API-Key": ""}) + assert exc.value.status_code == 500 From 1a7ca04cc5b908c37b2b78efd1cf683bfdd2163b Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 02:29:26 +0000 Subject: [PATCH 049/187] fix(proxy): carry litellm_call_id through endpoint specific error logs and failure responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/anthropic_endpoints/endpoints.py | 10 ++- litellm/proxy/batches_endpoints/endpoints.py | 27 ++++--- litellm/proxy/common_request_processing.py | 41 +++++++--- .../common_utils/openai_error_payload.py | 6 ++ litellm/proxy/image_endpoints/endpoints.py | 24 +++--- .../pass_through_endpoints.py | 12 ++- litellm/proxy/proxy_server.py | 46 ++++++----- litellm/proxy/rerank_endpoints/endpoints.py | 20 +++-- litellm/proxy/utils.py | 9 ++- .../anthropic_endpoints/test_endpoints.py | 81 +++++++++++++++++++ .../proxy/batches_endpoints/test_endpoints.py | 23 ++++++ .../common_utils/test_openai_error_payload.py | 6 ++ .../proxy/image_endpoints/test_endpoints.py | 68 ++++++++++++++++ .../test_pass_through_endpoints.py | 39 +++++++++ .../proxy/rerank_endpoints/test_endpoints.py | 51 ++++++++++-- .../proxy/test_common_request_processing.py | 8 +- tests/test_litellm/proxy/test_proxy_server.py | 43 ++++++++++ tests/test_litellm/proxy/test_proxy_utils.py | 31 +++++++ .../proxy/utils/helpers/test_error_helpers.py | 15 ++++ 19 files changed, 485 insertions(+), 75 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index d4cb3b84ee4..f673a82654d 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -8,7 +8,6 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse import litellm -from litellm._logging import verbose_proxy_logger from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.anthropic.experimental_pass_through.context_management import ( @@ -22,7 +21,9 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, create_response, + log_llm_api_exception, proxy_exception_from_http_exception, + resolve_litellm_call_id, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.openai_error_payload import ( @@ -218,7 +219,7 @@ async def anthropic_response( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=base_llm_response_processor.data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e) + log_llm_api_exception(e, base_llm_response_processor.litellm_call_id) if isinstance(e, ProxyException): return _anthropic_error_json_response(e, request) @@ -231,7 +232,7 @@ async def anthropic_response( # Get headers headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, - call_id=data.get("litellm_call_id", ""), + call_id=base_llm_response_processor.litellm_call_id, model_id=model_id, version=version, response_cost=0, @@ -288,6 +289,7 @@ async def count_tokens( """ from litellm.proxy.proxy_server import token_counter as internal_token_counter + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: request_data: Final = await _read_request_body(request=request) data: Final[dict] = {**request_data} @@ -339,7 +341,7 @@ async def count_tokens( detail=detail, ) except Exception as e: - verbose_proxy_logger.exception("litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - %s", e) + log_llm_api_exception(e, litellm_call_id) raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 5c4bacd757c..e3767d06e7d 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -17,7 +17,11 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + request_litellm_call_id, +) from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.common_utils.openai_endpoint_utils import ( @@ -383,8 +387,9 @@ async def create_batch( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e) - raise handle_exception_on_proxy(e) + litellm_call_id: Final = request_litellm_call_id(data) + log_llm_api_exception(e, litellm_call_id) + raise handle_exception_on_proxy(e, litellm_call_id) @router.get( @@ -674,8 +679,9 @@ async def retrieve_batch( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e) - raise handle_exception_on_proxy(e) + litellm_call_id: Final = request_litellm_call_id(data) + log_llm_api_exception(e, litellm_call_id) + raise handle_exception_on_proxy(e, litellm_call_id) @router.get( @@ -725,6 +731,7 @@ async def list_batches( ) verbose_proxy_logger.debug("GET /v1/batches after=%s limit=%s", after, limit) + data: dict = {} try: if llm_router is None: raise HTTPException( @@ -856,8 +863,9 @@ async def list_batches( original_exception=e, request_data={"after": after, "limit": limit}, ) - verbose_proxy_logger.error("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e) - raise handle_exception_on_proxy(e) + litellm_call_id: Final = request_litellm_call_id(data) + log_llm_api_exception(e, litellm_call_id) + raise handle_exception_on_proxy(e, litellm_call_id) @router.post( @@ -1079,8 +1087,9 @@ async def cancel_batch( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e) - raise handle_exception_on_proxy(e) + litellm_call_id: Final = request_litellm_call_id(data) + log_llm_api_exception(e, litellm_call_id) + raise handle_exception_on_proxy(e, litellm_call_id) ###################################################################### diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 46b222a4fc9..ba8f48b88be 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -7,7 +7,18 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, TypeAlias, TypeVar, overload +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NamedTuple, + Protocol, + TypeAlias, + TypeVar, + overload, + runtime_checkable, +) import anyio import httpx @@ -1452,7 +1463,19 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool: _CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request" -def _log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None: +@runtime_checkable +class _CarriesLitellmCallId(Protocol): + litellm_call_id: str | None + + +def request_litellm_call_id(data: Mapping[str, object]) -> str | None: + logging_obj: Final = data.get("litellm_logging_obj") + logged_id: Final = logging_obj.litellm_call_id if isinstance(logging_obj, _CarriesLitellmCallId) else None + call_id: Final = logged_id or data.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None + + +def log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None: if getattr(e, "status_code", None) == 499 and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL: verbose_proxy_logger.info( "litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, " @@ -1532,6 +1555,10 @@ class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict): self.data = data + @property + def litellm_call_id(self) -> str | None: + return request_litellm_call_id(self.data) + @staticmethod def _merge_passthrough_streaming_headers( response_headers: httpx.Headers | dict | None, @@ -3429,11 +3456,7 @@ class ProxyBaseLLMRequestProcessing: version: str | None = None, ): """Raises ProxyException (OpenAI API compatible) if an exception is raised""" - logging_obj: Final[LiteLLMLoggingObj | None] = self.data.get("litellm_logging_obj", None) - _log_llm_api_exception( - e, - (logging_obj.litellm_call_id if logging_obj is not None else None) or self.data.get("litellm_call_id"), - ) + log_llm_api_exception(e, self.litellm_call_id) # Allow callbacks to transform the error response transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -3463,9 +3486,7 @@ class ProxyBaseLLMRequestProcessing: custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, - call_id=( - _litellm_logging_obj.litellm_call_id if _litellm_logging_obj else self.data.get("litellm_call_id") - ), + call_id=self.litellm_call_id, model_id=model_id, version=version, response_cost=0, diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index fe23ab2c4b6..d4312d93559 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -52,3 +52,9 @@ def openai_error_param(exc: object) -> str | None: serializes as JSON ``null``.""" carried: Final = attribute_of(exc, "param") return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None + + +def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None: # mutable-ok: ProxyException.headers + if litellm_call_id is None: + return None + return {"x-litellm-call-id": litellm_call_id} # mutable-ok: ProxyException mutates its headers dict diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 3f044855ce8..30406bbcaae 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -1,6 +1,5 @@ import asyncio import io -import traceback from collections.abc import Sequence from typing import Final, get_type_hints @@ -9,19 +8,23 @@ from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, from fastapi.responses import ORJSONResponse import litellm -from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + resolve_litellm_call_id, +) from litellm.proxy.common_utils.http_parsing_utils import ( coerce_numeric_form_fields, numeric_form_fields, ) from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, + litellm_call_id_headers, openai_error_param, openai_error_type, ) @@ -92,6 +95,7 @@ async def image_generation( ) data = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() @@ -106,6 +110,7 @@ async def image_generation( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id if isinstance(model, str): reject_url_valued_destination("model", model) @@ -153,9 +158,7 @@ async def image_generation( response = await llm_call ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) ### CALL HOOKS ### - modify outgoing data (guardrails, otel, etc.) response = await proxy_logging_obj.post_call_success_hook( @@ -168,7 +171,7 @@ async def image_generation( cache_key: Final = hidden_params.get("cache_key", None) or "" api_base: Final = hidden_params.get("api_base", None) or "" response_cost: Final = hidden_params.get("response_cost", None) or "" - litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or "" + response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" fastapi_response.headers.update( ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -179,7 +182,7 @@ async def image_generation( version=version, response_cost=response_cost, model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - call_id=litellm_call_id, + call_id=response_call_id, request_data=data, hidden_params=hidden_params, ) @@ -200,13 +203,13 @@ async def image_generation( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error("litellm.proxy.proxy_server.image_generation(): Exception occured - %s", e) - verbose_proxy_logger.debug(traceback.format_exc()) + log_llm_api_exception(e, litellm_call_id) if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), param=openai_error_param(e), + headers=litellm_call_id_headers(litellm_call_id), code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: @@ -215,6 +218,7 @@ async def image_generation( message=getattr(e, "message", error_msg), type=openai_error_type(e, error_status_code(e, 500)), param=openai_error_param(e), + headers=litellm_call_id_headers(litellm_call_id), openai_code=getattr(e, "code", None), code=error_status_code(e, 500), ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 686544d352c..e93b1232836 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -72,7 +72,9 @@ from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_end from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, + log_llm_api_exception, open_sse_before_first_byte, + resolve_litellm_call_id, ) from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, @@ -80,6 +82,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( ) from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, + litellm_call_id_headers, openai_error_param, openai_error_type, ) @@ -197,6 +200,7 @@ async def chat_completion_pass_through_endpoint( ) data = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: body: Final = await request.body() body_str: Final = body.decode() @@ -224,6 +228,7 @@ async def chat_completion_pass_through_endpoint( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id # override with user settings, these are params passed via cli if user_temperature: @@ -290,9 +295,7 @@ async def chat_completion_pass_through_endpoint( response_cost: Final = hidden_params.get("response_cost", None) or "" ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) verbose_proxy_logger.debug("final response: %s", response) @@ -313,12 +316,13 @@ async def chat_completion_pass_through_endpoint( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - %s", e) + log_llm_api_exception(e, litellm_call_id) error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=openai_error_type(e, error_status_code(e, 500)), param=openai_error_param(e), + headers=litellm_call_id_headers(litellm_call_id), code=error_status_code(e, 500), ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7964556531..de42dcec8d0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -348,7 +348,10 @@ from litellm.proxy.common_request_processing import ( _is_azure_model_router_request, _should_return_raw_model_name, create_response, + log_llm_api_exception, open_sse_before_first_byte, + request_litellm_call_id, + resolve_litellm_call_id, ttft_keepalive_interval, ) from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( @@ -387,6 +390,7 @@ from litellm.proxy.common_utils.model_listing_utils import ( from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) +from litellm.proxy.common_utils.openai_error_payload import litellm_call_id_headers from litellm.proxy.common_utils.periodic_reload_schedule import ( MODEL_COST_MAP_RELOAD_PARAM_NAME, clear_reload_interval, @@ -11291,12 +11295,14 @@ async def completion( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - %s", e) + litellm_call_id: Final = request_litellm_call_id(data) + log_llm_api_exception(e, litellm_call_id) error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), openai_code=getattr(e, "code", None), code=getattr(e, "status_code", 500), ) @@ -11454,6 +11460,7 @@ async def moderations( """ global proxy_logging_obj data: dict = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() @@ -11468,6 +11475,7 @@ async def moderations( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id data["model"] = ( general_settings.get("moderation_model", None) # server default @@ -11494,9 +11502,7 @@ async def moderations( response: Final = await llm_call ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) ### RESPONSE HEADERS ### hidden_params: Final = getattr(response, "_hidden_params", {}) or {} @@ -11522,7 +11528,7 @@ async def moderations( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.moderations(): Exception occured - %s", e) + log_llm_api_exception(e, litellm_call_id) if isinstance(e, ProxyException): raise if isinstance(e, HTTPException): @@ -11530,6 +11536,7 @@ async def moderations( message=getattr(e, "message", str(e)), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: @@ -11538,6 +11545,7 @@ async def moderations( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), code=getattr(e, "status_code", 500), ) @@ -11576,6 +11584,7 @@ async def audio_speech( """ global proxy_logging_obj data: dict = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() @@ -11590,6 +11599,7 @@ async def audio_speech( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id if data.get("user", None) is None and user_api_key_dict.user_id is not None: data["user"] = user_api_key_dict.user_id @@ -11612,9 +11622,7 @@ async def audio_speech( response: Final = await llm_call ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) ### RESPONSE HEADERS ### hidden_params: Final = getattr(response, "_hidden_params", {}) or {} @@ -11622,7 +11630,7 @@ async def audio_speech( cache_key: Final = hidden_params.get("cache_key", None) or "" api_base: Final = hidden_params.get("api_base", None) or "" response_cost: Final = hidden_params.get("response_cost", None) or "" - litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or "" + response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, @@ -11633,7 +11641,7 @@ async def audio_speech( response_cost=response_cost, model_region=getattr(user_api_key_dict, "allowed_model_region", ""), fastest_response_batch_completion=None, - call_id=litellm_call_id, + call_id=response_call_id, request_data=data, hidden_params=hidden_params, ) @@ -11669,14 +11677,14 @@ async def audio_speech( original_exception=e, request_data=data, ) - verbose_proxy_logger.error("litellm.proxy.proxy_server.audio_speech(): Exception occured - %s", e) - verbose_proxy_logger.debug(traceback.format_exc()) + log_llm_api_exception(e, litellm_call_id) if isinstance(e, (ProxyException, HTTPException)): raise e raise ProxyException( message=getattr(e, "message", f"{e}"), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), openai_code=getattr(e, "code", None), code=getattr(e, "status_code", 500), ) @@ -11705,6 +11713,7 @@ async def audio_transcriptions( """ global proxy_logging_obj data: dict = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: # Use orjson to parse JSON data, orjson speeds up requests significantly form_data: Final = await get_form_data(request) @@ -11719,6 +11728,7 @@ async def audio_transcriptions( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id if data.get("user", None) is None and user_api_key_dict.user_id is not None: data["user"] = user_api_key_dict.user_id @@ -11775,9 +11785,7 @@ async def audio_transcriptions( file_object.close() # close the file read in by io library ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) ### RESPONSE HEADERS ### hidden_params: Final = getattr(response, "_hidden_params", {}) or {} @@ -11785,7 +11793,7 @@ async def audio_transcriptions( cache_key: Final = hidden_params.get("cache_key", None) or "" api_base: Final = hidden_params.get("api_base", None) or "" response_cost: Final = hidden_params.get("response_cost", None) or "" - litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or "" + response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" additional_headers: Final[dict] = hidden_params.get("additional_headers", {}) or {} fastapi_response.headers.update( @@ -11797,7 +11805,7 @@ async def audio_transcriptions( version=version, response_cost=response_cost, model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - call_id=litellm_call_id, + call_id=response_call_id, request_data=data, hidden_params=hidden_params, **additional_headers, @@ -11819,12 +11827,13 @@ async def audio_transcriptions( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.exception("litellm.proxy.proxy_server.audio_transcription(): Exception occured - %s", e) + log_llm_api_exception(e, litellm_call_id) if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) else: @@ -11833,6 +11842,7 @@ async def audio_transcriptions( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), + headers=litellm_call_id_headers(litellm_call_id), openai_code=getattr(e, "code", None), code=getattr(e, "status_code", 500), ) diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index 16cd7368e4a..4f5eb411e44 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -7,12 +7,16 @@ import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import ORJSONResponse -from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + resolve_litellm_call_id, +) from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, + litellm_call_id_headers, openai_error_param, openai_error_type, ) @@ -55,6 +59,7 @@ async def rerank( ) data = {} + litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) try: body: Final = await request.body() data = orjson.loads(body) @@ -68,6 +73,7 @@ async def rerank( version=version, proxy_config=proxy_config, ) + data["litellm_call_id"] = litellm_call_id ### CALL HOOKS ### - modify incoming data / reject request before calling the model data = await proxy_logging_obj.pre_call_hook(user_api_key_dict=user_api_key_dict, data=data, call_type="rerank") @@ -82,9 +88,7 @@ async def rerank( response: Final = await llm_call ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") - ) + asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) ### RESPONSE HEADERS ### hidden_params: Final = getattr(response, "_hidden_params", {}) or {} @@ -95,7 +99,7 @@ async def rerank( fastapi_response.headers.update( ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, - call_id=hidden_params.get("litellm_call_id", None) or data.get("litellm_call_id", None), + call_id=hidden_params.get("litellm_call_id", None) or litellm_call_id, model_id=model_id, cache_key=cache_key, api_base=api_base, @@ -113,12 +117,13 @@ async def rerank( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error("litellm.proxy.proxy_server.rerank(): Exception occured - %s", e) + log_llm_api_exception(e, litellm_call_id) if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), param=openai_error_param(e), + headers=litellm_call_id_headers(litellm_call_id), code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: @@ -127,5 +132,6 @@ async def rerank( message=getattr(e, "message", error_msg), type=openai_error_type(e, error_status_code(e, 500)), param=openai_error_param(e), + headers=litellm_call_id_headers(litellm_call_id), code=error_status_code(e, 500), ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 479bd0a55af..53782227998 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -38,7 +38,7 @@ from litellm.proxy._types import ( SpendLogsMetadata, SpendLogsPayload, ) -from litellm.proxy.common_utils.openai_error_payload import openai_error_param +from litellm.proxy.common_utils.openai_error_payload import litellm_call_id_headers, openai_error_param from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse @@ -3031,7 +3031,7 @@ class ProxyLogging: if litellm_logging_obj is None: from litellm._uuid import uuid - request_data["litellm_call_id"] = str(uuid.uuid4()) + request_data.setdefault("litellm_call_id", str(uuid.uuid4())) user_api_key_logged_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( user_api_key_dict=user_api_key_dict ) @@ -7638,7 +7638,7 @@ def _recreate_writer_on_read_only_transaction(prisma_client: "PrismaClient | Non asyncio.create_task(prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction")) -def handle_exception_on_proxy(e: Exception) -> ProxyException: +def handle_exception_on_proxy(e: Exception, litellm_call_id: str | None = None) -> ProxyException: """ Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible """ @@ -7650,11 +7650,13 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: _recreate_writer_on_read_only_transaction(prisma_client) + headers: Final = litellm_call_id_headers(litellm_call_id) if isinstance(e, HTTPException): return ProxyException( message=getattr(e, "detail", f"error({e})"), type=ProxyErrorTypes.internal_server_error, param=openai_error_param(e), + headers=headers, code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), ) elif isinstance(e, ProxyException): @@ -7664,6 +7666,7 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: message=str(e), type=ProxyErrorTypes.internal_server_error, param=openai_error_param(e), + headers=headers, code=_status_code, ) diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index f809fadc879..e4b15cfdcd3 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -3,12 +3,14 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objec """ import json +import logging import unittest from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +from litellm._logging import verbose_proxy_logger from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -285,6 +287,85 @@ class TestFailureHookRequestData: assert hook_request_data["litellm_logging_obj"] == "logging-obj-sentinel" +class TestErrorLogCarriesCallId: + """LIT-7836: the /v1/messages and /v1/messages/count_tokens error lines must carry + the request's litellm_call_id, rendered in the message and as a structured field.""" + + @pytest.fixture(autouse=True) + def propagating_proxy_logger(self): + verbose_proxy_logger.propagate = True + try: + yield + finally: + verbose_proxy_logger.propagate = False + + @staticmethod + def _error_record(caplog: pytest.LogCaptureFixture) -> logging.LogRecord: + return next(r for r in caplog.records if "Exception occured" in r.getMessage()) + + @pytest.mark.asyncio + async def test_messages_failure_log_carries_call_id(self, caplog: pytest.LogCaptureFixture): + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + call_id = "messages-call-7836" + + async def fake_process(self, **kwargs): + self.data = {**self.data, "litellm_call_id": call_id} + raise RuntimeError("provider timeout") + + request = MagicMock() + request.headers = {} + + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the provider failure happens inside this call; the test targets the endpoint's except block + patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), + ): + mock_logging.post_call_failure_hook = AsyncMock() + response = await ep.anthropic_response( + fastapi_response=MagicMock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert response.status_code == 500 + record = self._error_record(caplog) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() + + @pytest.mark.asyncio + async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture): + from fastapi import HTTPException + + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + + call_id = "count-tokens-call-7836" + request = MagicMock() + request.headers = {"x-litellm-call-id": call_id} + + with ( + patch.object( # test-quality-ok: endpoint reads the body via a module function; no injection seam + ep, + "_read_request_body", + new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}), + ), + patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block + caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), + pytest.raises(HTTPException) as raised, + ): + await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + + assert raised.value.status_code == 500 + record = self._error_record(caplog) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a37c8ff2bb4..cf805b384d2 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -31,6 +31,7 @@ cannot drift without a test failure. import base64 import json +import logging from contextlib import ExitStack from dataclasses import dataclass from typing import Any, Dict, Optional @@ -1088,6 +1089,28 @@ async def test_create__exception_calls_failure_hook(harness, openai_env_creds): assert harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom" +async def test_create__exception_carries_the_litellm_call_id(harness, openai_env_creds, caplog): + call_id = "lit7836-batch-call-id" + set_body( + harness, + { + "input_file_id": "file-plain", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "litellm_call_id": call_id, + }, + ) + harness.litellm_acreate.side_effect = ValueError("provider boom") + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised: + await call_create(harness) + + assert raised.value.headers["x-litellm-call-id"] == call_id + record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() + + # =========================================================================== # # # # GET /v1/batches/{batch_id} - retrieve_batch routing-contract tests # diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index 90850840ab4..df775916046 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -6,6 +6,7 @@ from fastapi import HTTPException from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, + litellm_call_id_headers, openai_error_param, openai_error_type, ) @@ -158,3 +159,8 @@ def test_a_stringified_none_type_or_param_is_treated_as_absent(): assert carried.type == "None" assert openai_error_type(carried, 400) == "invalid_request_error" assert openai_error_param(carried) is None + + +def test_a_failed_request_answers_with_the_call_id_it_was_logged_under(): + assert litellm_call_id_headers("call-7836") == {"x-litellm-call-id": "call-7836"} + assert litellm_call_id_headers(None) is None diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py index d8b3eef98bd..d03832bf6d0 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py @@ -1,5 +1,7 @@ import asyncio import copy +import logging +from collections.abc import Iterator from types import SimpleNamespace from typing import Any, Dict @@ -10,6 +12,7 @@ from fastapi.testclient import TestClient from starlette.requests import Request from starlette.responses import Response +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.image_endpoints import endpoints @@ -211,3 +214,68 @@ async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(mon await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()) assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "404") + + +@pytest.fixture +def propagating_proxy_logger() -> Iterator[None]: + verbose_proxy_logger.propagate = True + try: + yield + finally: + verbose_proxy_logger.propagate = False + + +@pytest.mark.asyncio +async def test_failure_log_carries_the_callers_litellm_call_id( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, propagating_proxy_logger: None +) -> None: + """LIT-7836: the /v1/images/generations error line must carry the litellm_call_id + the client sent, both rendered in the message and as a structured record field.""" + call_id = "images-call-7836" + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + async def fake_pre_call_hook(*, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> dict[str, object]: + return data + + async def fake_post_call_failure_hook(**_: object) -> None: + return None + + async def failing_route_request(**_: object) -> None: + raise HTTPException(status_code=401, detail={"error": "invalid api key"}) + + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", + SimpleNamespace(pre_call_hook=fake_pre_call_hook, post_call_failure_hook=fake_post_call_failure_hook), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version") + monkeypatch.setattr("litellm.proxy.image_endpoints.endpoints.route_request", failing_route_request) + + body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk"}) + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/images/generations", + "headers": [(b"x-litellm-call-id", call_id.encode())], + }, + receive, + ) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised: + await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()) + + assert raised.value.headers["x-litellm-call-id"] == call_id + record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 0fc961cf8c9..f5348c8adc1 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -6021,3 +6021,42 @@ async def test_chat_completion_pass_through_endpoint_answers_an_openai_typed_err ) assert (raised.value.type, raised.value.param, raised.value.code) == ("invalid_request_error", None, "400") + + +@pytest.mark.asyncio +async def test_chat_completion_pass_through_endpoint_failure_carries_the_callers_litellm_call_id( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + call_id = "lit7836-pass-through-call-id" + proxy_logging = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + proxy_logging.post_call_failure_hook = AsyncMock() + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging) + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + request = MagicMock(spec=Request) + request.headers = Headers({"x-litellm-call-id": call_id}) + request.body = AsyncMock( + return_value=json.dumps({"model": "unknown-model", "messages": [{"role": "user", "content": "hi"}]}).encode() + ) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised: + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert raised.value.headers["x-litellm-call-id"] == call_id + record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() diff --git a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py index ea858e04e0f..52d12dd1813 100644 --- a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py @@ -3,6 +3,8 @@ Tests for rerank_endpoints/endpoints.py response headers. """ import json +import logging +from collections.abc import Iterator from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -10,6 +12,7 @@ from fastapi import HTTPException, Request, Response import litellm.proxy.common_request_processing as common_request_processing_mod import litellm.proxy.proxy_server as proxy_server_mod +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.rerank_endpoints.endpoints import rerank from litellm.types.utils import RerankResponse @@ -28,7 +31,7 @@ HIDDEN_PARAMS = { } -def _build_request() -> Request: +def _build_request(headers: tuple[tuple[bytes, bytes], ...] = ()) -> Request: body = json.dumps({"model": "rerank-model", "query": "q", "documents": ["a", "b"]}).encode() async def receive(): @@ -39,7 +42,7 @@ def _build_request() -> Request: "type": "http", "method": "POST", "path": "/rerank", - "headers": [(b"content-type", b"application/json")], + "headers": [(b"content-type", b"application/json"), *headers], "query_string": b"", }, receive=receive, @@ -56,7 +59,7 @@ async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response: proxy_logging_obj.update_request_status = AsyncMock() async def fake_add_litellm_data_to_request(**kwargs): - return {**kwargs["data"], "litellm_call_id": "call-123"} + return dict(kwargs["data"]) async def fake_route_request(**kwargs): async def _call(): @@ -72,7 +75,7 @@ async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response: patch.object(proxy_server_mod, "version", "1.2.3"), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler ): await rerank( - request=_build_request(), + request=_build_request(headers=((b"x-litellm-call-id", b"call-123"),)), fastapi_response=fastapi_response, user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), ) @@ -121,7 +124,11 @@ async def test_rerank_omits_detailed_timing_headers_when_disabled(): async def _rerank_failure( - failure: Exception, *, raised_before_routing: bool, monkeypatch: pytest.MonkeyPatch + failure: Exception, + *, + raised_before_routing: bool, + monkeypatch: pytest.MonkeyPatch, + headers: tuple[tuple[bytes, bytes], ...] = (), ) -> ProxyException: proxy_logging_obj = MagicMock() proxy_logging_obj.pre_call_hook = AsyncMock( @@ -143,13 +150,45 @@ async def _rerank_failure( with pytest.raises(ProxyException) as raised: await rerank( - request=_build_request(), + request=_build_request(headers), fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), ) return raised.value +@pytest.fixture +def propagating_proxy_logger() -> Iterator[None]: + verbose_proxy_logger.propagate = True + try: + yield + finally: + verbose_proxy_logger.propagate = False + + +@pytest.mark.asyncio +async def test_failure_log_carries_the_callers_litellm_call_id( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, propagating_proxy_logger: None +) -> None: + """LIT-7836: the /rerank error line must carry the same litellm_call_id the client + sent, both in the rendered message and as a structured log record field.""" + call_id = "rerank-call-7836" + failure = HTTPException(status_code=401, detail={"error": "invalid api key"}) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + raised = await _rerank_failure( + failure, + raised_before_routing=False, + monkeypatch=monkeypatch, + headers=((b"x-litellm-call-id", call_id.encode()),), + ) + + assert raised.headers["x-litellm-call-id"] == call_id + record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() + + @pytest.mark.asyncio async def test_a_model_the_router_cannot_serve_answers_an_openai_typed_error(monkeypatch: pytest.MonkeyPatch): """A bare HTTPException carries no type or param, so the tail used to ship the diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 099204cd6c6..a621a97d448 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -8212,7 +8212,7 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_ """Regression for LIT-6043: expected 4xx errors log without formatting a traceback; unexpected errors keep logger.exception behavior.""" from litellm._logging import verbose_proxy_logger - from litellm.proxy.common_request_processing import _log_llm_api_exception + from litellm.proxy.common_request_processing import log_llm_api_exception verbose_proxy_logger.propagate = True try: @@ -8220,7 +8220,7 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_ try: raise exc except Exception as raised: - _log_llm_api_exception(raised, "call-id-for-traceback-test") + log_llm_api_exception(raised, "call-id-for-traceback-test") finally: verbose_proxy_logger.propagate = False @@ -8778,14 +8778,14 @@ class TestErrorLogCarriesCallId: from litellm._logging import verbose_proxy_logger from litellm.proxy.common_request_processing import ( _CLIENT_DISCONNECT_DETAIL, - _log_llm_api_exception, + log_llm_api_exception, ) call_id: Final = str(uuid.uuid4()) verbose_proxy_logger.propagate = True try: with caplog.at_level("INFO", logger="LiteLLM Proxy"): - _log_llm_api_exception( + log_llm_api_exception( HTTPException(status_code=499, detail=_CLIENT_DISCONNECT_DETAIL), call_id, ) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 6f55449abab..85c1a23c473 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2,6 +2,7 @@ import asyncio import contextlib import importlib import json +import logging import os import re import socket @@ -12927,6 +12928,48 @@ async def test_moderations_response_carries_litellm_call_id_header(): assert fastapi_response.headers["x-litellm-model-id"] == "mod-deployment-1" +@pytest.mark.asyncio +async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplog): + """LIT-7836: the /v1/moderations error line must carry the litellm_call_id the + client sent, rendered in the message and as a structured log record field.""" + from litellm._logging import verbose_proxy_logger + from litellm.proxy._types import ProxyException + + call_id = "moderations-call-7836" + + async def passthrough_add_litellm_data(data, **kwargs): + return data + + request = MagicMock() + request.headers = {"x-litellm-call-id": call_id} + request.body = AsyncMock(return_value=b'{"input": "hi"}') + fake_logging = MagicMock() + fake_logging.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + fake_logging.post_call_failure_hook = AsyncMock() + + verbose_proxy_logger.propagate = True + try: + with ( + patch.object(proxy_server_module, "add_litellm_data_to_request", new=passthrough_add_litellm_data), # test-quality-ok: the route reads this module global, no injection point + patch.object(proxy_server_module, "route_request", new=AsyncMock(side_effect=Exception("bad key"))), # test-quality-ok: fakes the provider failure so the real route's error log is observable + patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), + pytest.raises(ProxyException) as raised, + ): + await proxy_server_module.moderations( + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0), + ) + finally: + verbose_proxy_logger.propagate = False + + assert raised.value.headers["x-litellm-call-id"] == call_id + record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) + assert record.litellm_call_id == call_id + assert call_id in record.getMessage() + + @pytest.mark.asyncio async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(monkeypatch): from litellm.proxy.agent_endpoints.agent_registry import ( diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 94ccc2762c5..df18e5c6093 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -160,6 +160,37 @@ async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params(): assert "litellm_metadata" not in captured["optional_params"] +@pytest.mark.asyncio +async def test_proxy_only_error_log_keeps_the_request_litellm_call_id(monkeypatch: pytest.MonkeyPatch): + """LIT-7836: a route that already stamped the caller's litellm_call_id must + keep it when the failure is a proxy-only error, so the spend-log row and the + error line share one id instead of a fresh uuid minted here.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + call_id: Final = "caller-supplied-7836" + captured: dict[str, object] = {} + + def fake_pre_call(self, *args, **kwargs): + captured["litellm_call_id"] = self.litellm_call_id + + async def _noop_async_failure(self, *args, **kwargs): + return None + + monkeypatch.setattr(Logging, "pre_call", fake_pre_call) + monkeypatch.setattr(Logging, "async_failure_handler", _noop_async_failure) + request_data: Final[dict[str, object]] = {"model": "gpt-4o", "input": "hi", "litellm_call_id": call_id} + + await ProxyLogging(user_api_key_cache=DualCache())._handle_logging_proxy_only_error( + request_data=request_data, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-bad", request_route="/v1/moderations"), + route="/v1/moderations", + original_exception=Exception("bad key"), + ) + + assert request_data["litellm_call_id"] == call_id + assert captured["litellm_call_id"] == call_id + + def test_get_model_group_info_order(): from litellm import Router from litellm.proxy.proxy_server import _get_model_group_info diff --git a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py index 117c5aa3081..278d11f95b0 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py +++ b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py @@ -176,6 +176,21 @@ def test_handle_exception_on_proxy_error_path_none_input_wraps_as_500(): } +@pytest.mark.parametrize( + "exc", + [HTTPException(status_code=401, detail="bad key"), ValueError("provider boom")], + ids=["http_exception", "generic_exception"], +) +def test_handle_exception_on_proxy_returns_the_litellm_call_id_header(exc: Exception): + result = handle_exception_on_proxy(exc, "call-7836") + + assert result.headers == {"x-litellm-call-id": "call-7836"} + + +def test_handle_exception_on_proxy_sends_no_call_id_header_when_the_request_has_none(): + assert handle_exception_on_proxy(ValueError("provider boom")).headers == {} + + @pytest.mark.asyncio async def test_handle_exception_on_proxy_read_only_transaction_forces_writer_recreate( monkeypatch: pytest.MonkeyPatch, From f12feed9a9f0f7008050b3a969e7eeaeec6744e0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 02:39:39 +0000 Subject: [PATCH 050/187] test(proxy): expect litellm_call_id in the image generation call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/proxy_unit_tests/test_proxy_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 9c8dd90dd2b..1fcdaa67143 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -809,6 +809,7 @@ def test_img_gen(mock_aimage_generation, client_no_auth): n=1, size="1024x1024", imageConfig={"aspectRatio": "9:16", "imageSize": "1K"}, + litellm_call_id=mock.ANY, metadata=mock.ANY, proxy_server_request=mock.ANY, secret_fields=mock.ANY, From 1aa2e19ee4dc46a861c65ba7ebb506d41c307908 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 15 Sep 2026 19:52:00 -0700 Subject: [PATCH 051/187] test(together_ai): stop depending on a serverless model we do not control Together moved openai/gpt-oss-20b off serverless, so three tests started failing with a 400 model_not_available from the live API. None of them was really testing Together: they cover provider-prefix parsing, prompt shaping and streaming, all litellm side. Mock the transport and assert those, so the tests answer to our code instead of a vendor catalog. --- tests/local_testing/test_completion.py | 109 ++++++++++++++++++------- 1 file changed, 80 insertions(+), 29 deletions(-) diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index 43ed57f63af..a6ac112a3b7 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -11,6 +11,7 @@ import io from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import litellm @@ -57,21 +58,49 @@ def test_response_model_none(): assert isinstance(x, litellm.ModelResponse) +TOGETHER_AI_CHAT_URL = "https://api.together.ai/v1/chat/completions" + + +def _together_ai_chat_response(content="Hello!"): + return httpx.Response( + 200, + json={ + "id": "chatcmpl-together", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-oss-20b", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + request=httpx.Request("POST", TOGETHER_AI_CHAT_URL), + ) + + def test_completion_custom_provider_model_name(): - try: - litellm.cache = None + litellm.cache = None + with patch.object( + HTTPHandler, "post", return_value=_together_ai_chat_response() + ) as mock_post: response = completion( model="together_ai/openai/gpt-oss-20b", messages=messages, logger_fn=logger_fn, + api_key="fake-key", ) - # Add assertions here to check the-response - print(response) - print(response["choices"][0]["finish_reason"]) - except litellm.Timeout as e: - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") + + assert mock_post.call_args.kwargs["url"] == TOGETHER_AI_CHAT_URL + assert json.loads(mock_post.call_args.kwargs["data"])["model"] == "openai/gpt-oss-20b" + assert response.choices[0].finish_reason == "stop" def _openai_mock_response(*args, **kwargs) -> litellm.ModelResponse: @@ -2804,12 +2833,11 @@ def test_completion_together_ai_llama(): # test_completion_together_ai() def test_customprompt_together_ai(): - try: - litellm.set_verbose = False - litellm.num_retries = 0 - print("in test_customprompt_together_ai") - print(litellm.success_callback) - print(litellm._async_success_callback) + litellm.set_verbose = False + litellm.num_retries = 0 + with patch.object( + HTTPHandler, "post", return_value=_together_ai_chat_response() + ) as mock_post: response = completion( model="together_ai/openai/gpt-oss-20b", messages=messages, @@ -2827,14 +2855,14 @@ def test_customprompt_together_ai(): "post_message": "<|im_end|>", }, }, + api_key="fake-key", ) - print(response) - except litellm.exceptions.Timeout as e: - print(f"Timeout Error") - pass - except Exception as e: - print(f"ERROR TYPE {type(e)}") - pytest.fail(f"Error occurred: {e}") + + body = json.loads(mock_post.call_args.kwargs["data"]) + assert body["messages"] == messages + assert "prompt" not in body + assert "roles" not in body + assert response.choices[0].finish_reason == "stop" # test_customprompt_together_ai() @@ -3648,19 +3676,42 @@ def test_completion_together_ai_stream(): litellm.set_verbose = True user_message = "Write 1pg about YC & litellm" messages = [{"content": user_message, "role": "user"}] - try: + sse_body = ( + 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' + '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{"role":"assistant",' + '"content":"YC"},"finish_reason":null}]}\n\n' + 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' + '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{"content":" and ' + 'litellm"},"finish_reason":null}]}\n\n' + 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' + '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{},' + '"finish_reason":"stop"}]}\n\n' + "data: [DONE]\n\n" + ) + stream_response = httpx.Response( + 200, + content=sse_body.encode(), + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", TOGETHER_AI_CHAT_URL), + ) + + with patch.object( + HTTPHandler, "post", return_value=stream_response + ) as mock_post: response = completion( model="together_ai/openai/gpt-oss-20b", messages=messages, stream=True, max_tokens=5, + api_key="fake-key", ) - print(response) - for chunk in response: - print(chunk) - # print(string_response) - except Exception as e: - pytest.fail(f"Error occurred: {e}") + chunks = list(response) + + assert json.loads(mock_post.call_args.kwargs["data"])["stream"] is True + assert "".join( + chunk.choices[0].delta.content or "" for chunk in chunks + ) == "YC and litellm" + assert chunks[-1].choices[0].finish_reason == "stop" # test_completion_together_ai_stream() From be506936bd9b64776a2730b8e4c4c55f4f000491 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:52:40 -0700 Subject: [PATCH 052/187] fix(mcp): preserve explicit caller authorization credentials --- .../_experimental/mcp_server/upstream.py | 11 +++---- .../mcp_server/test_mcp_server_manager.py | 31 +++++++++++++++++++ 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index 49aca5cf31b..2a9312e17a1 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -36,17 +36,16 @@ def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> b return True -def validate_static_credential( - server: MCPServer, headers: Mapping[str, str], *, header_slot: str | None = None, openapi: bool = False -) -> Result[None, CredError]: +def validate_static_credential(server: MCPServer, headers: Mapping[str, str]) -> Result[None, CredError]: if server.auth_type not in _STATIC_MODES or server.transport == MCPTransport.stdio: return Ok(None) default_slot: Final = "X-API-Key" if server.auth_type == MCPAuth.api_key else "Authorization" slots: Final = frozenset( name.lower() for name in ( - header_slot or server.upstream_token_header or default_slot, - "Authorization" if openapi else default_slot, + server.upstream_token_header or default_slot, + default_slot, + "Authorization", ) ) values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) @@ -75,7 +74,7 @@ def validate_openapi_credentials( headers: Final = merge_openapi_headers( server.static_headers or {}, forwarded_headers, caller_authorization, resolved_headers ) - match validate_static_credential(server, headers, openapi=True): + match validate_static_credential(server, headers): case Error(error): raise_public(error) case Ok(): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index c8dba1c1554..a582632259e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13693,3 +13693,34 @@ class TestProtectedCredentialPreparation: with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Custom": "", "X-API-Key": ""}) assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("custom_slot", [None, "X-Custom"]) + @pytest.mark.parametrize("source", ["caller", "forwarded"]) + async def test_api_key_preserves_explicit_authorization_credential( + self, custom_slot: str | None, source: str + ) -> None: + server: Final = MCPServer( + server_id="caller-auth", name="caller-auth", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, upstream_token_header=custom_slot, + ) + headers: Final = {"Authorization": "Bearer caller-credential", "X-API-Key": ""} + client: Final = await MCPServerManager()._create_mcp_client( + server, mcp_auth_header=headers if source == "caller" else None, + extra_headers=headers if source == "forwarded" else None, + ) + request: Final = await client.prepare_request_auth() + assert request.headers["Authorization"] == "Bearer caller-credential" + assert request.headers["X-API-Key"] == "" + assert custom_slot is None or custom_slot not in request.headers + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["", " ", "Bearer", "Basic", "token", "ApiKey"]) + async def test_api_key_rejects_authorization_without_a_credential(self, value: str) -> None: + server: Final = MCPServer( + server_id="caller-empty", name="caller-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header={"Authorization": value}) + assert exc.value.status_code == 500 From 581c613f6680b1daf0e1da140250987d61ede4ac Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 03:02:50 +0000 Subject: [PATCH 053/187] fix(proxy): keep litellm_call_id on shaped errors and list_batches failure hook Already shaped ProxyException and HTTPException errors passing through the moderations, audio speech, Anthropic Messages, and handle_exception_on_proxy paths now answer with the x-litellm-call-id header the route logged under, without overwriting a header the exception was raised with. The GET /v1/batches failure hook receives the resolved request data so the spend log request_id matches the response header and the error log Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/anthropic_endpoints/endpoints.py | 5 +- litellm/proxy/batches_endpoints/endpoints.py | 2 +- .../common_utils/openai_error_payload.py | 12 +++- litellm/proxy/proxy_server.py | 18 +++-- litellm/proxy/utils.py | 8 ++- .../anthropic_endpoints/test_endpoints.py | 30 ++++++++ .../proxy/batches_endpoints/test_endpoints.py | 18 +++++ .../common_utils/test_openai_error_payload.py | 25 +++++++ tests/test_litellm/proxy/test_proxy_server.py | 71 ++++++++++++++++++- .../proxy/utils/helpers/test_error_helpers.py | 8 ++- 10 files changed, 184 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index f673a82654d..644778bcb9f 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -30,6 +30,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( error_status_code, openai_error_param, openai_error_type, + with_litellm_call_id, ) from litellm.types.utils import TokenCountResponse @@ -222,7 +223,9 @@ async def anthropic_response( log_llm_api_exception(e, base_llm_response_processor.litellm_call_id) if isinstance(e, ProxyException): - return _anthropic_error_json_response(e, request) + return _anthropic_error_json_response( + with_litellm_call_id(e, base_llm_response_processor.litellm_call_id), request + ) # Extract model_id from request metadata (same as success path) litellm_metadata: Final = data.get("litellm_metadata", {}) or {} diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index e3767d06e7d..f37c06aea97 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -861,7 +861,7 @@ async def list_batches( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, - request_data={"after": after, "limit": limit}, + request_data={**data, "after": after, "limit": limit}, ) litellm_call_id: Final = request_litellm_call_id(data) log_llm_api_exception(e, litellm_call_id) diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index d4312d93559..cbc8c78d4f9 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -9,6 +9,9 @@ from typing import Final from fastapi import status from litellm.constants import STRINGIFIED_NONE +from litellm.proxy._types import ProxyException + +LITELLM_CALL_ID_HEADER: Final = "x-litellm-call-id" _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( { @@ -57,4 +60,11 @@ def openai_error_param(exc: object) -> str | None: def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None: # mutable-ok: ProxyException.headers if litellm_call_id is None: return None - return {"x-litellm-call-id": litellm_call_id} # mutable-ok: ProxyException mutates its headers dict + return {LITELLM_CALL_ID_HEADER: litellm_call_id} # mutable-ok: ProxyException mutates its headers dict + + +def with_litellm_call_id(exc: ProxyException, litellm_call_id: str | None) -> ProxyException: + """The same error object, answering with ``x-litellm-call-id`` when it was raised without one.""" + if litellm_call_id is not None: + exc.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id) + return exc diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index de42dcec8d0..03a425fedb2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -390,7 +390,11 @@ from litellm.proxy.common_utils.model_listing_utils import ( from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) -from litellm.proxy.common_utils.openai_error_payload import litellm_call_id_headers +from litellm.proxy.common_utils.openai_error_payload import ( + LITELLM_CALL_ID_HEADER, + litellm_call_id_headers, + with_litellm_call_id, +) from litellm.proxy.common_utils.periodic_reload_schedule import ( MODEL_COST_MAP_RELOAD_PARAM_NAME, clear_reload_interval, @@ -11530,7 +11534,7 @@ async def moderations( ) log_llm_api_exception(e, litellm_call_id) if isinstance(e, ProxyException): - raise + raise with_litellm_call_id(e, litellm_call_id) if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), @@ -11678,8 +11682,14 @@ async def audio_speech( request_data=data, ) log_llm_api_exception(e, litellm_call_id) - if isinstance(e, (ProxyException, HTTPException)): - raise e + if isinstance(e, ProxyException): + raise with_litellm_call_id(e, litellm_call_id) + if isinstance(e, HTTPException): + raise HTTPException( + status_code=e.status_code, + detail=e.detail, + headers={LITELLM_CALL_ID_HEADER: litellm_call_id, **(e.headers or {})}, + ) raise ProxyException( message=getattr(e, "message", f"{e}"), type=getattr(e, "type", "None"), diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 53782227998..1376020a907 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -38,7 +38,11 @@ from litellm.proxy._types import ( SpendLogsMetadata, SpendLogsPayload, ) -from litellm.proxy.common_utils.openai_error_payload import litellm_call_id_headers, openai_error_param +from litellm.proxy.common_utils.openai_error_payload import ( + litellm_call_id_headers, + openai_error_param, + with_litellm_call_id, +) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse @@ -7660,7 +7664,7 @@ def handle_exception_on_proxy(e: Exception, litellm_call_id: str | None = None) code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), ) elif isinstance(e, ProxyException): - return e + return with_litellm_call_id(e, litellm_call_id) _status_code: Final = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) return ProxyException( message=str(e), diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index e4b15cfdcd3..9a9ccd9a213 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -336,6 +336,36 @@ class TestErrorLogCarriesCallId: assert record.litellm_call_id == call_id assert call_id in record.getMessage() + @pytest.mark.asyncio + async def test_messages_already_shaped_failure_answers_with_the_call_id(self): + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth + + call_id = "messages-call-7836-shaped" + + async def fake_process(self, **kwargs): + self.data = {**self.data, "litellm_call_id": call_id} + raise ProxyException(message="budget exceeded", type=ProxyErrorTypes.budget_exceeded, param="key", code=402) + + request = MagicMock() + request.headers = {} + + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})), # test-quality-ok: endpoint reads the body via a module function; no injection seam + patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process), # test-quality-ok: the proxy shaped failure happens inside this call; the test targets the endpoint's except block + patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam + ): + mock_logging.post_call_failure_hook = AsyncMock() + response = await ep.anthropic_response( + fastapi_response=MagicMock(), + request=request, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert response.status_code == 402 + assert response.headers["x-litellm-call-id"] == call_id + @pytest.mark.asyncio async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture): from fastapi import HTTPException diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index cf805b384d2..d9bfb3fe3da 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -1976,6 +1976,24 @@ async def test_list__exception_calls_failure_hook(list_harness): assert list_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom" +@pytest.mark.asyncio +async def test_list__failure_hook_and_response_share_the_request_litellm_call_id(list_harness): + call_id = "lit7836-list-batches-call-id" + list_harness.pre_call.side_effect = lambda **kw: ( + {**list_harness.body["body"], "litellm_call_id": call_id}, + MagicMock(), + ) + list_harness.litellm_alist.side_effect = ValueError("provider boom") + + with pytest.raises(ProxyException) as raised: + await call_list(list_harness, after="batch-0", limit=5) + + failure_request_data = list_harness.logging.post_call_failure_hook.call_args.kwargs["request_data"] + assert failure_request_data["litellm_call_id"] == call_id + assert (failure_request_data["after"], failure_request_data["limit"]) == ("batch-0", 5) + assert raised.value.headers["x-litellm-call-id"] == call_id + + # =========================================================================== # # # # POST /v1/batches/{batch_id}/cancel - cancel_batch routing-contract tests # diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index df775916046..c09b8742b50 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -9,6 +9,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( litellm_call_id_headers, openai_error_param, openai_error_type, + with_litellm_call_id, ) @@ -164,3 +165,27 @@ def test_a_stringified_none_type_or_param_is_treated_as_absent(): def test_a_failed_request_answers_with_the_call_id_it_was_logged_under(): assert litellm_call_id_headers("call-7836") == {"x-litellm-call-id": "call-7836"} assert litellm_call_id_headers(None) is None + + +def test_an_already_shaped_proxy_error_answers_with_the_call_id_it_was_logged_under(): + raised_without_id = ProxyException(message="budget exceeded", type="budget_exceeded", param="key", code=402) + + carried = with_litellm_call_id(raised_without_id, "call-7836") + + assert carried is raised_without_id + assert carried.headers == {"x-litellm-call-id": "call-7836"} + assert (carried.message, carried.type, carried.param, carried.code) == ( + "budget exceeded", + "budget_exceeded", + "key", + "402", + ) + + +def test_a_proxy_error_keeps_the_call_id_it_was_raised_with(): + raised_with_id = ProxyException( + message="nope", type="None", param=None, code=400, headers={"x-litellm-call-id": "first"} + ) + + assert with_litellm_call_id(raised_with_id, "second").headers == {"x-litellm-call-id": "first"} + assert with_litellm_call_id(ProxyException(message="nope", type="None", param=None, code=400), None).headers == {} diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 85c1a23c473..f3c477e4ad0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -20,7 +20,7 @@ import fastapi.routing import httpx import pytest import yaml -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException from fastapi.encoders import jsonable_encoder from fastapi.staticfiles import StaticFiles from fastapi.testclient import TestClient @@ -32,7 +32,7 @@ from litellm.caching.caching import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache -from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, ProxyErrorTypes, ProxyException, TokenCountRequest, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, initialize @@ -12970,6 +12970,73 @@ async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplo assert call_id in record.getMessage() +@pytest.mark.asyncio +async def test_moderations_already_shaped_failure_answers_with_the_callers_litellm_call_id(): + """LIT-7836: a ProxyException raised inside /v1/moderations is re-raised unwrapped but still + answers with the caller's x-litellm-call-id so the client can join it to the error log.""" + call_id = "moderations-call-7836-shaped" + exc = ProxyException(message="budget exceeded", type=ProxyErrorTypes.budget_exceeded, param="key", code=402) + + request = MagicMock() + request.headers = {"x-litellm-call-id": call_id} + request.body = AsyncMock(return_value=b'{"input": "hi"}') + fake_logging = MagicMock() + fake_logging.post_call_failure_hook = AsyncMock() + + with ( + patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point + patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + pytest.raises(ProxyException) as raised, + ): + await proxy_server_module.moderations( + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0), + ) + + assert raised.value is exc + assert raised.value.code == "402" + assert raised.value.headers["x-litellm-call-id"] == call_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "exc", + [ + HTTPException(status_code=401, detail="bad key"), + ProxyException(message="budget exceeded", type=ProxyErrorTypes.budget_exceeded, param="key", code=402), + ], + ids=["http_exception", "proxy_exception"], +) +async def test_audio_speech_already_shaped_failure_answers_with_the_callers_litellm_call_id(exc: Exception): + """LIT-7836: /v1/audio/speech re-raises HTTP and proxy shaped failures unchanged, and they must + still answer with the caller's x-litellm-call-id.""" + call_id = "speech-call-7836-shaped" + + request = MagicMock() + request.headers = {"x-litellm-call-id": call_id} + request.body = AsyncMock(return_value=b'{"model": "tts-1", "input": "hi", "voice": "alloy"}') + fake_logging = MagicMock() + fake_logging.post_call_failure_hook = AsyncMock() + + with ( + patch.object(proxy_server_module, "add_litellm_data_to_request", new=AsyncMock(side_effect=exc)), # test-quality-ok: the route reads this module global, no injection point + patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + pytest.raises(type(exc)) as raised, + ): + await proxy_server_module.audio_speech( + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0), + ) + + if isinstance(exc, HTTPException): + assert (raised.value.status_code, raised.value.detail) == (401, "bad key") + else: + assert raised.value is exc + assert raised.value.headers["x-litellm-call-id"] == call_id + + @pytest.mark.asyncio async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(monkeypatch): from litellm.proxy.agent_endpoints.agent_registry import ( diff --git a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py index 278d11f95b0..df399c8b1d2 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py +++ b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py @@ -178,8 +178,12 @@ def test_handle_exception_on_proxy_error_path_none_input_wraps_as_500(): @pytest.mark.parametrize( "exc", - [HTTPException(status_code=401, detail="bad key"), ValueError("provider boom")], - ids=["http_exception", "generic_exception"], + [ + HTTPException(status_code=401, detail="bad key"), + ValueError("provider boom"), + ProxyException(message="already wrapped", type=ProxyErrorTypes.budget_exceeded.value, param="key", code=402), + ], + ids=["http_exception", "generic_exception", "already_proxy_exception"], ) def test_handle_exception_on_proxy_returns_the_litellm_call_id_header(exc: Exception): result = handle_exception_on_proxy(exc, "call-7836") From 8c046e13bdcbd61f570f10f674ff57b2dcf19afb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 15 Sep 2026 20:17:42 -0700 Subject: [PATCH 054/187] test(together_ai): move request-shape checks to the mapped file, drop the live ones Together moved openai/gpt-oss-20b off serverless and three tests in test_completion.py died on a live 400. None of them needed Together to be up: streaming is already covered live by tests/e2e/llm_translation/test_together_ai_e2e.py, which picks its model from the cost map instead of pinning one, and the other two are request-shape questions. Delete all three and assert the two shapes in the mapped transformation file: the provider prefix is stripped without eating the rest of a slashed model name, and custom role wrappers never reach the request. --- tests/local_testing/test_completion.py | 125 ------------------ .../test_together_ai_chat_transformation.py | 64 +++++++++ 2 files changed, 64 insertions(+), 125 deletions(-) diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index a6ac112a3b7..25c6c50251d 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -11,7 +11,6 @@ import io from unittest.mock import AsyncMock, MagicMock, patch -import httpx import pytest import litellm @@ -58,51 +57,6 @@ def test_response_model_none(): assert isinstance(x, litellm.ModelResponse) -TOGETHER_AI_CHAT_URL = "https://api.together.ai/v1/chat/completions" - - -def _together_ai_chat_response(content="Hello!"): - return httpx.Response( - 200, - json={ - "id": "chatcmpl-together", - "object": "chat.completion", - "created": 1, - "model": "openai/gpt-oss-20b", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": content}, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 1, - "completion_tokens": 1, - "total_tokens": 2, - }, - }, - request=httpx.Request("POST", TOGETHER_AI_CHAT_URL), - ) - - -def test_completion_custom_provider_model_name(): - litellm.cache = None - with patch.object( - HTTPHandler, "post", return_value=_together_ai_chat_response() - ) as mock_post: - response = completion( - model="together_ai/openai/gpt-oss-20b", - messages=messages, - logger_fn=logger_fn, - api_key="fake-key", - ) - - assert mock_post.call_args.kwargs["url"] == TOGETHER_AI_CHAT_URL - assert json.loads(mock_post.call_args.kwargs["data"])["model"] == "openai/gpt-oss-20b" - assert response.choices[0].finish_reason == "stop" - - def _openai_mock_response(*args, **kwargs) -> litellm.ModelResponse: new_response = MagicMock() new_response.headers = {"hello": "world"} @@ -2832,40 +2786,6 @@ def test_completion_together_ai_llama(): # test_completion_together_ai() -def test_customprompt_together_ai(): - litellm.set_verbose = False - litellm.num_retries = 0 - with patch.object( - HTTPHandler, "post", return_value=_together_ai_chat_response() - ) as mock_post: - response = completion( - model="together_ai/openai/gpt-oss-20b", - messages=messages, - roles={ - "system": { - "pre_message": "<|im_start|>system\n", - "post_message": "<|im_end|>", - }, - "assistant": { - "pre_message": "<|im_start|>assistant\n", - "post_message": "<|im_end|>", - }, - "user": { - "pre_message": "<|im_start|>user\n", - "post_message": "<|im_end|>", - }, - }, - api_key="fake-key", - ) - - body = json.loads(mock_post.call_args.kwargs["data"]) - assert body["messages"] == messages - assert "prompt" not in body - assert "roles" not in body - assert response.choices[0].finish_reason == "stop" - - -# test_customprompt_together_ai() def response_format_tests(response: litellm.ModelResponse): @@ -3672,51 +3592,6 @@ async def test_acompletion_stream_watsonx(): # test_maritalk() -def test_completion_together_ai_stream(): - litellm.set_verbose = True - user_message = "Write 1pg about YC & litellm" - messages = [{"content": user_message, "role": "user"}] - sse_body = ( - 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' - '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{"role":"assistant",' - '"content":"YC"},"finish_reason":null}]}\n\n' - 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' - '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{"content":" and ' - 'litellm"},"finish_reason":null}]}\n\n' - 'data: {"id":"chatcmpl-together","object":"chat.completion.chunk","created":1,' - '"model":"openai/gpt-oss-20b","choices":[{"index":0,"delta":{},' - '"finish_reason":"stop"}]}\n\n' - "data: [DONE]\n\n" - ) - stream_response = httpx.Response( - 200, - content=sse_body.encode(), - headers={"content-type": "text/event-stream"}, - request=httpx.Request("POST", TOGETHER_AI_CHAT_URL), - ) - - with patch.object( - HTTPHandler, "post", return_value=stream_response - ) as mock_post: - response = completion( - model="together_ai/openai/gpt-oss-20b", - messages=messages, - stream=True, - max_tokens=5, - api_key="fake-key", - ) - chunks = list(response) - - assert json.loads(mock_post.call_args.kwargs["data"])["stream"] is True - assert "".join( - chunk.choices[0].delta.content or "" for chunk in chunks - ) == "YC and litellm" - assert chunks[-1].choices[0].finish_reason == "stop" - - -# test_completion_together_ai_stream() - - def test_moderation(): response = litellm.moderation(input="i'm ishaan cto of litellm") print(response) diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py index 7eb7dc41d4f..a7347edb2c7 100644 --- a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py +++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py @@ -1108,3 +1108,67 @@ def test_get_optional_params_preserves_max_for_declared_levels_model(): ) assert optional_params["reasoning_effort"] == "max" + + +def _together_chat_transport() -> tuple[HTTPHandler, list[httpx.Request]]: + captured_requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-together", + "object": "chat.completion", + "created": 1234567890, + "model": TOOL_CALLING_MODEL, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + return client, captured_requests + + +def test_only_the_provider_prefix_is_stripped_from_a_slashed_model_name(): + client, captured_requests = _together_chat_transport() + + litellm.completion( + model=f"together_ai/{TOOL_CALLING_MODEL}", + messages=[{"role": "user", "content": "Hello!"}], + api_key="fake-key", + client=client, + ) + + assert "/" in TOOL_CALLING_MODEL + assert str(captured_requests[0].url) == "https://api.together.ai/v1/chat/completions" + assert json.loads(captured_requests[0].content)["model"] == TOOL_CALLING_MODEL + + +def test_custom_role_wrappers_never_reach_the_request(): + client, captured_requests = _together_chat_transport() + messages = [{"role": "user", "content": "Hello!"}] + + litellm.completion( + model=f"together_ai/{TOOL_CALLING_MODEL}", + messages=messages, + roles={ + "system": {"pre_message": "<|im_start|>system\n", "post_message": "<|im_end|>"}, + "assistant": {"pre_message": "<|im_start|>assistant\n", "post_message": "<|im_end|>"}, + "user": {"pre_message": "<|im_start|>user\n", "post_message": "<|im_end|>"}, + }, + api_key="fake-key", + client=client, + ) + + request_body = json.loads(captured_requests[0].content) + assert request_body["messages"] == messages + assert "prompt" not in request_body + assert "roles" not in request_body From 258176de762939aca13375abef5a474b4595e8d9 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 20:32:25 -0700 Subject: [PATCH 055/187] fix(mcp): validate rendered static credential payloads --- .../_experimental/mcp_server/upstream.py | 8 ++- .../mcp_server/test_mcp_server_manager.py | 65 ++++++++++++++++++- 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py index 2a9312e17a1..66840db21ce 100644 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ b/litellm/proxy/_experimental/mcp_server/upstream.py @@ -4,7 +4,7 @@ import base64 from collections.abc import Mapping from typing import Final -from litellm.experimental_mcp_client.client import MCPClient +from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import raise_public from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError @@ -24,13 +24,17 @@ def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> b return True if value.lower() in ("bearer", "basic", "token", "apikey"): return False + if auth_type in (MCPAuth.bearer_token, MCPAuth.token): + scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token" + credential: Final = strip_auth_scheme(value, scheme).strip() + return bool(credential) and credential.lower() != scheme.lower() if auth_type == MCPAuth.basic: parts: Final = value.split(None, 1) if len(parts) != 2 or parts[0].lower() != "basic": return False try: decoded: Final = base64.b64decode(parts[1], validate=True).strip() - return bool(decoded) and decoded.lower() != b"basic" + return b":" in decoded except ValueError: return False return True diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index a582632259e..54add273c24 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13639,7 +13639,7 @@ class TestProtectedCredentialPreparation: assert exc.value.status_code == 500 @pytest.mark.asyncio - @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc", "Basic QmFzaWM="]) + @pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc", "Basic QmFzaWM=", "Basic bm8tY29sb24="]) async def test_basic_headers_without_usable_credentials_reject(self, header: str) -> None: server = MCPServer(server_id="bad-basic", name="bad-basic", url="https://upstream.example/mcp", transport=MCPTransport.http, auth_type=MCPAuth.basic) @@ -13724,3 +13724,66 @@ class TestProtectedCredentialPreparation: with pytest.raises(HTTPException) as exc: await MCPServerManager()._create_mcp_client(server, mcp_auth_header={"Authorization": value}) assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["no-colon", "Basic bm8tY29sb24="]) + @pytest.mark.parametrize("source", ["configured", "caller"]) + async def test_basic_requires_a_username_password_separator(self, value: str, source: str) -> None: + server: Final = MCPServer( + server_id="basic-pair", name="basic-pair", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic, + authentication_token=value if source == "configured" else None, + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None) + assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["user:pass", "user:", ":pass", ":"]) + async def test_basic_preserves_username_password_pairs(self, value: str) -> None: + import base64 + + server: Final = MCPServer( + server_id="basic-valid", name="basic-valid", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.basic, authentication_token=value, + ) + client: Final = await MCPServerManager()._create_mcp_client(server) + request: Final = await client.prepare_request_auth() + scheme, encoded = request.headers["Authorization"].split(" ", 1) + assert scheme == "Basic" + assert base64.b64decode(encoded) == value.encode() + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,value", [ + (MCPAuth.bearer_token, "Bearer"), (MCPAuth.bearer_token, "Bearer "), (MCPAuth.bearer_token, "bearer"), + (MCPAuth.token, "token"), (MCPAuth.token, "token "), (MCPAuth.token, "TOKEN"), + ]) + @pytest.mark.parametrize("source", ["configured", "caller"]) + async def test_static_scheme_only_input_cannot_hide_behind_rendered_prefix( + self, auth_type: MCPAuthType, value: str, source: str + ) -> None: + server: Final = MCPServer( + server_id="empty-scheme", name="empty-scheme", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, + authentication_token=value if source == "configured" else None, + ) + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, mcp_auth_header=value if source == "caller" else None) + assert exc.value.status_code == 500 + + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,value,expected", [ + (MCPAuth.bearer_token, "token", "Bearer token"), + (MCPAuth.bearer_token, "Bearertoken", "Bearer Bearertoken"), + (MCPAuth.token, "tokenish", "token tokenish"), + ]) + async def test_static_credentials_that_resemble_schemes_remain_usable( + self, auth_type: MCPAuthType, value: str, expected: str + ) -> None: + server: Final = MCPServer( + server_id="real-token", name="real-token", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=value, + ) + client: Final = await MCPServerManager()._create_mcp_client(server) + request: Final = await client.prepare_request_auth() + assert request.headers["Authorization"] == expected From 7bce15f7d8493096ced51023f7a239091bfcc1f2 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 02:45:43 +0000 Subject: [PATCH 056/187] fix(otel): propagate W3C trace context on passthrough upstream requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 36 ++++++++ .../pass_through_endpoints.py | 30 +++--- .../otel/test_otel_v2_components.py | 81 ++++++++++++++++- .../test_pass_through_endpoints.py | 91 +++++++++++++++++++ 4 files changed, 224 insertions(+), 14 deletions(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 21e61c71fb7..851097a17e7 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -310,6 +310,42 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None: return _PROPAGATOR.extract(carrier) +def _outgoing_trace_context(inbound_headers: Mapping[str, str] | None = None) -> Context | None: + root: Final = request_root_span() + if root is not None: + return context_from_span(root) + + current: Final = get_current() + if is_recordable_span(get_current_span(current)): + return current + + if inbound_headers is None: + return None + inbound_context: Final = extract_traceparent(inbound_headers) + if inbound_context is None or not is_recordable_span(get_current_span(inbound_context)): + return None + return inbound_context + + +def inject_trace_context( + headers: Mapping[str, str], + inbound_headers: Mapping[str, str] | None = None, +) -> dict[str, str]: + """``headers`` plus W3C ``traceparent``/``tracestate`` for the current request's span. + + Parent preference: the anchored request root span, then the ambient active span, + then the trace context the caller sent inbound. Only trace context is injected, + never Baggage, so per-request identity baggage cannot leak upstream. Unchanged + when no valid span context exists anywhere. + """ + context: Final = _outgoing_trace_context(inbound_headers) + if context is None: + return dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier + carrier: Final = dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier + _PROPAGATOR.inject(carrier, context=context) + return carrier + + # The OTLP destinations this request's key or team pointed its traces at, resolved # once during auth. A ``ContextVar`` for the same reason the root span above is one: # it rides the request task's context into the ``asyncio.create_task`` children that diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 686544d352c..1c4f8d59c9c 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -955,6 +955,7 @@ async def pass_through_request( general_settings.pass_through_request_timeout, then 600s. """ from litellm.exceptions import ModifyResponseException + from litellm.integrations.otel.plumbing.context import inject_trace_context from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.pass_through_endpoints.passthrough_guardrails import ( PassthroughGuardrailHandler, @@ -985,6 +986,7 @@ async def pass_through_request( headers=headers, forward_headers=forward_headers, ) + headers = inject_trace_context(headers, inbound_headers=_safe_get_request_headers(request)) requested_query_params: dict | None = query_params or dict(request.query_params) @@ -2199,20 +2201,22 @@ async def websocket_passthrough_request( await websocket.accept() verbose_proxy_logger.debug("WebSocket passthrough (%s): WebSocket connection accepted", endpoint) - # Prepare headers for the upstream connection - upstream_headers: Final = custom_headers.copy() + from litellm.integrations.otel.plumbing.context import inject_trace_context - if forward_headers: - # Forward relevant headers from the incoming request - incoming_headers: Final = dict(websocket.headers) - for header_name, header_value in incoming_headers.items(): - # Only forward certain headers to avoid conflicts - if header_name.lower() in [ - "authorization", - "x-api-key", - "x-goog-user-project", - ]: - upstream_headers[header_name] = header_value + incoming_headers: Final = dict(websocket.headers) # mutable-ok: websocket headers are copied for context extraction + forwarded_headers: Final = { # mutable-ok: assembled as the upstream header carrier + **custom_headers, + **{ + header_name: header_value + for header_name, header_value in incoming_headers.items() + if forward_headers + and header_name.lower() in frozenset(("authorization", "x-api-key", "x-goog-user-project")) + }, + } + upstream_headers: Final = inject_trace_context( + forwarded_headers, + inbound_headers=incoming_headers, + ) # Initialize logging object similar to HTTP passthrough team_callbacks: Final = _resolve_team_callback_wiring( diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index ae41c74944d..b1e65da1661 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -5,6 +5,7 @@ builders, and the registry validator's failure paths. Needs the OTel SDK.""" import json import threading from collections.abc import Iterator +from contextvars import Context as ContextVarContext from dataclasses import replace from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer @@ -15,6 +16,8 @@ pytest.importorskip("opentelemetry") from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( # noqa: E402 ExportTraceServiceRequest, ) +from opentelemetry import baggage # noqa: E402 +from opentelemetry.context import attach, detach # noqa: E402 from opentelemetry.sdk.metrics import MeterProvider # noqa: E402 from opentelemetry.sdk.metrics.export import InMemoryMetricReader # noqa: E402 from opentelemetry.sdk.trace import TracerProvider # noqa: E402 @@ -26,7 +29,10 @@ from opentelemetry.sdk.trace.export import ( # noqa: E402 from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( # noqa: E402 InMemorySpanExporter, ) -from opentelemetry.trace import SpanKind # noqa: E402 +from opentelemetry.trace import SpanKind, get_current_span # noqa: E402 +from opentelemetry.trace.propagation.tracecontext import ( # noqa: E402 + TraceContextTextMapPropagator, +) from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 @@ -464,6 +470,79 @@ def test_extract_traceparent(): assert ctx_mod.extract_traceparent({"x": "y"}) is None +def _test_tracer(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return provider.get_tracer("test") + + +def test_inject_trace_context_prefers_request_root_span(): + def run(): + tracer = _test_tracer() + with tracer.start_as_current_span("root") as root: + ctx_mod.set_request_root_span(root) + result = ctx_mod.inject_trace_context( + {"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"} + ) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return result, root, propagated + + result, root, propagated = ContextVarContext().run(run) + assert result["traceparent"] != "00-11111111111111111111111111111111-2222222222222222-01" + assert propagated.get_span_context().trace_id == root.get_span_context().trace_id + assert propagated.get_span_context().span_id == root.get_span_context().span_id + + +def test_inject_trace_context_uses_ambient_span_without_request_root(): + def run(): + tracer = _test_tracer() + with tracer.start_as_current_span("ambient") as ambient: + result = ctx_mod.inject_trace_context({}) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return ambient, propagated + + ambient, propagated = ContextVarContext().run(run) + assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id + assert propagated.get_span_context().span_id == ambient.get_span_context().span_id + + +def test_inject_trace_context_forwards_valid_inbound_context_without_span(): + inbound = {"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"} + + def run(): + result = ctx_mod.inject_trace_context({}, inbound_headers=inbound) + return get_current_span(TraceContextTextMapPropagator().extract(result)) + + propagated = ContextVarContext().run(run) + assert propagated.get_span_context().trace_id == int("0af7651916cd43dd8448eb211c80319c", 16) + assert propagated.get_span_context().span_id == int("b7ad6b7169203331", 16) + + +def test_inject_trace_context_returns_headers_unchanged_without_context(): + headers = {"x-custom": "value"} + + result = ContextVarContext().run(lambda: ctx_mod.inject_trace_context(headers)) + + assert result == headers + assert "traceparent" not in result + assert result is not headers + + +def test_inject_trace_context_does_not_forward_baggage(): + def run(): + tracer = _test_tracer() + with tracer.start_as_current_span("ambient"): + token = attach(baggage.set_baggage("litellm.team.id", "team")) + try: + return ctx_mod.inject_trace_context({}) + finally: + detach(token) + + result = ContextVarContext().run(run) + assert "baggage" not in result + + def test_set_request_baggage_empty_returns_context(): assert ctx_mod.set_request_baggage({}) is not None diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 0fc961cf8c9..c3ad900cc9e 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4270,6 +4270,40 @@ def _relay_client_request(method="GET"): return mock_request +@pytest.mark.asyncio +async def test_pass_through_request_propagates_active_trace_context(): + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.trace import get_current_span + from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + from litellm.proxy._types import UserAPIKeyAuth + + captured: dict[str, httpx.Headers] = {} + + def transport_handler(upstream_request: httpx.Request) -> httpx.Response: + captured["headers"] = upstream_request.headers + return httpx.Response(200, json={"ok": True}, request=upstream_request) + + fake_client, cleanup = _inject_fake_passthrough_client(httpx.MockTransport(transport_handler), timeout=None) + try: + with ExitStack() as stack: + _enter_relay_logging_mocks(stack, {}) + tracer = TracerProvider().get_tracer("test") + with tracer.start_as_current_span("passthrough") as span: + response = await pass_through_request( + request=_relay_client_request(method="POST"), + target="http://internal-api.test/v1/generate", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + finally: + cleanup() + await fake_client.aclose() + + assert response.status_code == 200 + propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) + assert propagated.get_span_context().trace_id == span.get_span_context().trace_id + + @pytest.mark.asyncio async def test_pass_through_request_relays_non_json_body_without_buffering(): """ @@ -4866,6 +4900,63 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): assert all(call.kwargs.get("code") != 1011 for call in websocket.close.await_args_list) +@pytest.mark.asyncio +async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch): + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.trace import get_current_span + from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + from starlette.websockets import WebSocketState + + captured: dict[str, dict[str, str]] = {} + upstream_ws = FakeUpstreamWebSocket(b"{}") + + def fake_connect(target, additional_headers): + captured["headers"] = additional_headers + return FakeUpstreamConnect(upstream_ws) + + websocket = MagicMock() + websocket.accept = AsyncMock() + websocket.send_text = AsyncMock() + websocket.send_bytes = AsyncMock() + websocket.receive = AsyncMock(return_value={"type": "websocket.disconnect"}) + websocket.close = AsyncMock() + websocket.headers = {} + websocket.client_state = WebSocketState.CONNECTED + websocket.application_state = WebSocketState.CONNECTED + tracer = TracerProvider().get_tracer("test") + + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_success_hook = AsyncMock() + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_worker = MagicMock() + mock_worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: async_coroutine.close() + ) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", + fake_connect, + ) + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER", + mock_worker, + ) + with tracer.start_as_current_span("websocket_passthrough") as span: + await websocket_passthrough_request( + websocket=websocket, + target="wss://upstream.example.test/v1/realtime", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/realtime", + accept_websocket=True, + ) + + propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) + assert propagated.get_span_context().trace_id == span.get_span_context().trace_id + + class ClosingUpstreamWebSocket: def __init__(self, close_exc: Exception): self._close_exc = close_exc From 7bcf848b39883c651514d8df0c366a8c89ce54e2 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 02:47:01 +0000 Subject: [PATCH 057/187] refactor(passthrough): hoist websocket forwarded header set Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 1c4f8d59c9c..186318dccc1 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2159,6 +2159,9 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None: return upstream_close +_WEBSOCKET_FORWARDED_HEADERS: Final = frozenset(("authorization", "x-api-key", "x-goog-user-project")) + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2183,6 +2186,7 @@ async def websocket_passthrough_request( cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream """ + from litellm.integrations.otel.plumbing.context import inject_trace_context from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -2201,22 +2205,16 @@ async def websocket_passthrough_request( await websocket.accept() verbose_proxy_logger.debug("WebSocket passthrough (%s): WebSocket connection accepted", endpoint) - from litellm.integrations.otel.plumbing.context import inject_trace_context - - incoming_headers: Final = dict(websocket.headers) # mutable-ok: websocket headers are copied for context extraction - forwarded_headers: Final = { # mutable-ok: assembled as the upstream header carrier + incoming_headers: Final = dict(websocket.headers) # mutable-ok: propagator carrier + forwarded_headers: Final = { # mutable-ok: propagator carrier **custom_headers, **{ header_name: header_value for header_name, header_value in incoming_headers.items() - if forward_headers - and header_name.lower() in frozenset(("authorization", "x-api-key", "x-goog-user-project")) + if forward_headers and header_name.lower() in _WEBSOCKET_FORWARDED_HEADERS }, } - upstream_headers: Final = inject_trace_context( - forwarded_headers, - inbound_headers=incoming_headers, - ) + upstream_headers: Final = inject_trace_context(forwarded_headers, inbound_headers=incoming_headers) # Initialize logging object similar to HTTP passthrough team_callbacks: Final = _resolve_team_callback_wiring( From d2e8b9c6565478c9ce1f5ba1cca834537a1759af Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 03:05:05 +0000 Subject: [PATCH 058/187] fix(otel): keep passthrough working when opentelemetry is not installed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints.py | 14 ++++++++++---- .../test_pass_through_endpoints.py | 11 +++++++++++ 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 186318dccc1..6d44b4fa6dc 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -955,7 +955,6 @@ async def pass_through_request( general_settings.pass_through_request_timeout, then 600s. """ from litellm.exceptions import ModifyResponseException - from litellm.integrations.otel.plumbing.context import inject_trace_context from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.pass_through_endpoints.passthrough_guardrails import ( PassthroughGuardrailHandler, @@ -986,7 +985,7 @@ async def pass_through_request( headers=headers, forward_headers=forward_headers, ) - headers = inject_trace_context(headers, inbound_headers=_safe_get_request_headers(request)) + headers = _with_trace_context(headers, inbound_headers=_safe_get_request_headers(request)) requested_query_params: dict | None = query_params or dict(request.query_params) @@ -2162,6 +2161,14 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None: _WEBSOCKET_FORWARDED_HEADERS: Final = frozenset(("authorization", "x-api-key", "x-goog-user-project")) +def _with_trace_context(headers: Mapping[str, str], inbound_headers: Mapping[str, str]) -> dict[str, str]: + try: + from litellm.integrations.otel.plumbing.context import inject_trace_context + except ImportError: + return dict(headers) # mutable-ok: matches inject_trace_context's carrier return type + return inject_trace_context(headers, inbound_headers=inbound_headers) + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2186,7 +2193,6 @@ async def websocket_passthrough_request( cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream """ - from litellm.integrations.otel.plumbing.context import inject_trace_context from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -2214,7 +2220,7 @@ async def websocket_passthrough_request( if forward_headers and header_name.lower() in _WEBSOCKET_FORWARDED_HEADERS }, } - upstream_headers: Final = inject_trace_context(forwarded_headers, inbound_headers=incoming_headers) + upstream_headers: Final = _with_trace_context(forwarded_headers, inbound_headers=incoming_headers) # Initialize logging object similar to HTTP passthrough team_callbacks: Final = _resolve_team_callback_wiring( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index c3ad900cc9e..05371dccb23 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -2,6 +2,7 @@ import asyncio import json import logging import os +import sys from collections.abc import Callable from contextlib import ExitStack, contextmanager from io import BytesIO @@ -29,6 +30,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( resolve_pass_through_request_timeout, resolve_llm_passthrough_timeout, websocket_passthrough_request, + _with_trace_context, ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -46,6 +48,15 @@ import litellm MESSAGE_START_SSE_FRAME = b'event: message_start\ndata: {"type": "message_start"}\n\n' +def test_with_trace_context_without_opentelemetry(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setitem(sys.modules, "litellm.integrations.otel.plumbing.context", None) + + headers = _with_trace_context({"authorization": "x"}, {}) + + assert headers == {"authorization": "x"} + assert "traceparent" not in headers + + # Test is_multipart def test_is_multipart(): # Test with multipart content type From e4a12510f647fb9ef5a790f93f1b829b78f7d905 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 03:22:10 +0000 Subject: [PATCH 059/187] test(otel): cover websocket forwarded headers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints/test_pass_through_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 05371dccb23..d81e3ae5af8 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4931,7 +4931,7 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch websocket.send_bytes = AsyncMock() websocket.receive = AsyncMock(return_value={"type": "websocket.disconnect"}) websocket.close = AsyncMock() - websocket.headers = {} + websocket.headers = {"authorization": "Bearer client"} websocket.client_state = WebSocketState.CONNECTED websocket.application_state = WebSocketState.CONNECTED tracer = TracerProvider().get_tracer("test") @@ -4959,7 +4959,7 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch target="wss://upstream.example.test/v1/realtime", custom_headers={}, user_api_key_dict=UserAPIKeyAuth(), - forward_headers=False, + forward_headers=True, endpoint="/realtime", accept_websocket=True, ) From 68cc12e848fb445b42aa27e6c005e91b17595e41 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 03:36:35 +0000 Subject: [PATCH 060/187] test(otel): assert websocket forwarded header reaches upstream Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/pass_through_endpoints/test_pass_through_endpoints.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d81e3ae5af8..a21f695493b 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4966,6 +4966,7 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) assert propagated.get_span_context().trace_id == span.get_span_context().trace_id + assert captured["headers"]["authorization"] == "Bearer client" class ClosingUpstreamWebSocket: From 7ba073aa2668820a35c69dede3ef46ed7b837cc6 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 11 Sep 2026 03:47:09 +0000 Subject: [PATCH 061/187] test(otel): cover websocket trace propagation with forwarding on and off Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints/test_pass_through_endpoints.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index a21f695493b..0603b5cb2ff 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4912,7 +4912,8 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): @pytest.mark.asyncio -async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch): +@pytest.mark.parametrize("forward_headers", [True, False]) +async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch, forward_headers: bool): from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import get_current_span from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator @@ -4959,14 +4960,14 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch target="wss://upstream.example.test/v1/realtime", custom_headers={}, user_api_key_dict=UserAPIKeyAuth(), - forward_headers=True, + forward_headers=forward_headers, endpoint="/realtime", accept_websocket=True, ) propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) assert propagated.get_span_context().trace_id == span.get_span_context().trace_id - assert captured["headers"]["authorization"] == "Bearer client" + assert captured["headers"].get("authorization") == ("Bearer client" if forward_headers else None) class ClosingUpstreamWebSocket: From 2b6184d76867fd38a990d2df124b7d1cd808ca6c Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:08:24 +0000 Subject: [PATCH 062/187] fix(proxy): resolve rate-limit fallbacks after model normalization and retry from a client-request snapshot The fallback retry in _pre_call_with_fallbacks re-entered common_processing_pre_call_logic with data already enriched by the first pass, so add_litellm_data_to_request deep-copied a metadata dict holding the live OTel span and the request failed with a 500 (cannot pickle '_thread.RLock') instead of the intended 429 or fallback. Capture the configured fallbacks and a snapshot of the client request before the first pass, look up the fallback chain by the normalized model group after the limiter raises, and run each fallback attempt on a fresh copy of that snapshot. Replaces the mock-heavy tests with a rig that runs the real v3 limiter and a live OTel span through the proxy_logging_obj seam Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 62 ++-- .../test_response_polling_pre_call_checks.py | 8 +- .../proxy/test_common_request_processing.py | 347 +++++++----------- 3 files changed, 166 insertions(+), 251 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3c8bb9b3c92..7c9a296965c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2066,20 +2066,12 @@ class ProxyBaseLLMRequestProcessing: ) -> tuple[dict, LiteLLMLoggingObj]: from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError - original_model: Final = self.data.get("model") - fallback_models: Final = ( - self._resolve_fallback_models( - model=original_model, - llm_router=llm_router, - user_api_key_dict=user_api_key_dict, - ) - if original_model - and isinstance(original_model, str) - and llm_router - and not self.data.get("disable_fallbacks") + configured_fallbacks: Final = ( + self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict) + if llm_router is not None and not self.data.get("disable_fallbacks") else None ) - pristine: Final = independent_snapshot(self.data) if fallback_models else None + pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None try: return await self.common_processing_pre_call_logic( @@ -2099,7 +2091,16 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, ) except ProxyRateLimitError as original_exc: - if not fallback_models or pristine is None: + rate_limited_data: Final = self.data + original_model: Final = rate_limited_data.get("model") + if pristine is None or not configured_fallbacks or not isinstance(original_model, str): + raise + + fallback_models: Final = self._resolve_fallback_models( + model=original_model, + fallbacks=configured_fallbacks, + ) + if not fallback_models: raise verbose_proxy_logger.info( @@ -2133,39 +2134,30 @@ class ProxyBaseLLMRequestProcessing: except ProxyRateLimitError: continue except BaseException: - self.data = pristine + self.data = rate_limited_data raise - self.data = pristine + self.data = rate_limited_data raise original_exc - def _resolve_fallback_models( - self, - model: str, - llm_router: Router, - user_api_key_dict: UserAPIKeyAuth, - ) -> list | None: - from litellm.router_utils.fallback_event_handlers import get_fallback_model_group - - fallbacks = None - + @staticmethod + def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None: key_router_settings: Final = user_api_key_dict.router_settings - if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings: - fallbacks = key_router_settings["fallbacks"] + key_fallbacks: Final = key_router_settings.get("fallbacks") if isinstance(key_router_settings, dict) else None + fallbacks: Final = key_fallbacks if key_fallbacks is not None else llm_router.fallbacks + return fallbacks if isinstance(fallbacks, list) and fallbacks else None - if fallbacks is None: - fallbacks = llm_router.fallbacks - - if not fallbacks: - return None + @staticmethod + def _resolve_fallback_models(model: str, fallbacks: list) -> list | None: + from litellm.router_utils.fallback_event_handlers import get_fallback_model_group fallback_model_group, generic_fallback_idx = get_fallback_model_group( fallbacks=fallbacks, model_group=model, ) - if fallback_model_group is None and generic_fallback_idx is not None: - fallback_model_group = fallbacks[generic_fallback_idx]["*"] - return fallback_model_group + if fallback_model_group is not None: + return fallback_model_group + return fallbacks[generic_fallback_idx]["*"] if generic_fallback_idx is not None else None @staticmethod def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str: diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py index 38f087f51ca..459834d0fd2 100644 --- a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py +++ b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py @@ -48,10 +48,10 @@ class TestSkipPreCallLogic: await processor.base_process_llm_request( request=MagicMock(spec=Request), fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth, router_settings=None), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), route_type="aresponses", proxy_logging_obj=mock_proxy_logging, - llm_router=MagicMock(fallbacks=None), + llm_router=MagicMock(), general_settings={}, proxy_config=MagicMock(), skip_pre_call_logic=True, @@ -87,10 +87,10 @@ class TestSkipPreCallLogic: await processor.base_process_llm_request( request=MagicMock(spec=Request), fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth, router_settings=None), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), route_type="aresponses", proxy_logging_obj=mock_proxy_logging, - llm_router=MagicMock(fallbacks=None), + llm_router=MagicMock(), general_settings={}, proxy_config=MagicMock(), ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 92af65b2637..f689dd62df6 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6365,247 +6365,170 @@ class TestPreCallWithFallbacksOnLocalRateLimit: call_type="acompletion", ) - @pytest.mark.asyncio - async def test_fallback_retries_from_pristine_request_data(self): - import threading + @staticmethod + def _v3_limiter_rig( + monkeypatch: pytest.MonkeyPatch, + user_api_key_dict: ProxyUserAPIKeyAuth, + fallbacks: list[dict[str, list[str]]], + ) -> tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]]: + """Real v3 limiter (the default ``parallel_request_limiter``) wired in through the + ``proxy_logging_obj`` seam, so ``common_processing_pre_call_logic`` runs for real: + ``add_litellm_data_to_request`` with a live OTel span, ``function_setup``, then the limiter.""" + from litellm.caching.caching import DualCache + from litellm.proxy import proxy_server + from litellm.proxy.hooks.parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + monkeypatch.setattr(proxy_server, "prisma_client", None) + limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) + limiter_models: list[str] = [] - primary_model = "gpt-4" - fallback_model = "gpt-3.5-turbo" + async def run_limiter(**kwargs): + limiter_models.append(kwargs["data"]["model"]) + await limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data=kwargs["data"], + call_type=kwargs["call_type"], + ) + return kwargs["data"] - processor = ProxyBaseLLMRequestProcessing( - data={ - "model": primary_model, - "messages": [{"role": "user", "content": "hi"}], - "metadata": {"tags": ["a"]}, - } + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=run_limiter) + router = litellm.Router( + model_list=[ + {"model_name": group, "litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "fake"}} + for chain in fallbacks + for group in (*chain.keys(), *(m for models in chain.values() for m in models)) + ], + fallbacks=fallbacks, ) + return proxy_logging_obj, router, proxy_server.ProxyConfig(), limiter_models - metadata_at_entry = [] - - async def mock_pre_call_logic(**kwargs): - copy.deepcopy(processor.data["metadata"]) - metadata_at_entry.append(dict(processor.data["metadata"])) - processor.data["metadata"]["litellm_parent_otel_span"] = threading.RLock() - processor.data["litellm_logging_obj"] = object() - if processor.data.get("model") == primary_model: - raise ProxyRateLimitError( - detail="TPM limit exceeded for gpt-4", - headers={"retry-after": "30"}, - ) - return processor.data, MagicMock() - - mock_router = MagicMock() - mock_router.fallbacks = [{primary_model: [fallback_model]}] - - with patch.object( - processor, - "common_processing_pre_call_logic", - side_effect=mock_pre_call_logic, - ): - data, logging_obj = await processor._pre_call_with_fallbacks( - request=MagicMock(), - general_settings={}, - proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), - version=None, - proxy_config=MagicMock(), - user_model=None, - user_temperature=None, - user_request_timeout=None, - user_max_tokens=None, - user_api_base=None, - model=primary_model, - route_type="acompletion", - llm_router=mock_router, - ) - - assert processor.data["model"] == fallback_model - assert metadata_at_entry[1] == {"tags": ["a"]} - - @pytest.mark.asyncio - async def test_exhausted_fallbacks_restore_pristine_request_data(self): - import threading - - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError - - primary_model = "gpt-4" - original_data = { - "model": primary_model, - "messages": [{"role": "user", "content": "hi"}], - "metadata": {"tags": ["a"]}, - } - processor = ProxyBaseLLMRequestProcessing(data=copy.deepcopy(original_data)) - - async def mock_pre_call_logic(**kwargs): - processor.data["metadata"]["litellm_parent_otel_span"] = threading.RLock() - processor.data["litellm_logging_obj"] = object() - raise ProxyRateLimitError( - detail=f"TPM limit exceeded for {processor.data.get('model')}", - headers={"retry-after": "30"}, - ) - - mock_router = MagicMock() - mock_router.fallbacks = [{primary_model: ["gpt-3.5-turbo"]}] - - with patch.object( - processor, - "common_processing_pre_call_logic", - side_effect=mock_pre_call_logic, - ): - with pytest.raises(ProxyRateLimitError, match="gpt-4"): - await processor._pre_call_with_fallbacks( - request=MagicMock(), - general_settings={}, - proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), - version=None, - proxy_config=MagicMock(), - user_model=None, - user_temperature=None, - user_request_timeout=None, - user_max_tokens=None, - user_api_base=None, - model=primary_model, - route_type="acompletion", - llm_router=mock_router, - ) - - assert processor.data == original_data - - @pytest.mark.asyncio - async def test_real_add_litellm_data_to_request_rerun_with_otel_span_falls_back(self): - from opentelemetry import trace + @staticmethod + def _otel_key(**limits) -> ProxyUserAPIKeyAuth: from opentelemetry.sdk.trace import TracerProvider - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError - from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request - from litellm.proxy.proxy_server import ProxyConfig + span = TracerProvider().get_tracer("test").start_span("proxy-request") + return ProxyUserAPIKeyAuth(api_key="hashed-key", parent_otel_span=span, **limits) - trace.set_tracer_provider(TracerProvider()) + @staticmethod + def _chat_request() -> Request: + return Request({"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []}) - primary_model = "gpt-4" - fallback_model = "gpt-3.5-turbo" - - request_mock = MagicMock(spec=Request) - request_mock.url = MagicMock() - request_mock.url.path = "/v1/chat/completions" - request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" - request_mock.method = "POST" - request_mock.query_params = {} - request_mock.headers = {"Content-Type": "application/json"} - request_mock.client = MagicMock() - request_mock.client.host = "127.0.0.1" - - user_api_key_dict = UserAPIKeyAuth( - parent_otel_span=trace.get_tracer("x").start_span("s"), - api_key="hashed-key", - user_id="u1", - team_id="t1", - metadata={}, - team_metadata={}, - team_member_tpm_limit=1000, + async def _pre_call( + self, + data: dict, + user_api_key_dict: ProxyUserAPIKeyAuth, + rig: tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]], + ) -> tuple[ProxyBaseLLMRequestProcessing, tuple[dict, object]]: + proxy_logging_obj, router, proxy_config, _ = rig + processor = ProxyBaseLLMRequestProcessing(data=data) + result = await processor._pre_call_with_fallbacks( + request=self._chat_request(), + general_settings={}, + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=user_api_key_dict, + version=None, + proxy_config=proxy_config, + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model=None, + route_type="acompletion", + llm_router=router, ) + return processor, result - processor = ProxyBaseLLMRequestProcessing( - data={ + @pytest.mark.asyncio + async def test_v3_limiter_with_otel_span_falls_back_from_client_request(self, monkeypatch: pytest.MonkeyPatch): + """Customer path: OTel on, per-key model RPM cap on the primary, a router fallback configured. + The first pass enriches ``data["metadata"]`` with the live span, then the limiter raises. The + fallback pass must start from the client's request again, so ``add_litellm_data_to_request`` + never deep-copies the span (the ``cannot pickle '_thread.RLock'`` 500).""" + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + key = self._otel_key(metadata={"model_rpm_limit": {primary_model: 1}}) + rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + + def client_request() -> dict: + return { "model": primary_model, "messages": [{"role": "user", "content": "hi"}], - "metadata": {"tags": ["a"]}, + "metadata": {"tags": ["client-tag"]}, } - ) - async def real_add_litellm_data_pre_call(**kwargs): - await add_litellm_data_to_request( - data=processor.data, - request=request_mock, - user_api_key_dict=user_api_key_dict, - proxy_config=ProxyConfig(), + _, (first_data, _) = await self._pre_call(client_request(), key, rig) + processor, (data, logging_obj) = await self._pre_call(client_request(), key, rig) + + assert first_data["model"] == primary_model + assert data["model"] == fallback_model + assert processor.data is data + assert data["litellm_logging_obj"] is logging_obj + assert logging_obj.model == fallback_model + requester_metadata = data["metadata"]["requester_metadata"] + assert requester_metadata["tags"] == ["client-tag"] + assert "litellm_parent_otel_span" not in requester_metadata + assert "user_api_key_auth" not in requester_metadata + assert data["metadata"]["litellm_parent_otel_span"] is key.parent_otel_span + assert rig[3] == [primary_model, primary_model, fallback_model] + + @pytest.mark.asyncio + async def test_v3_limiter_with_otel_span_returns_429_when_fallbacks_exhausted( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + key = self._otel_key(rpm_limit=1) + rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + + await self._pre_call(dict(request), key, rig) + processor = ProxyBaseLLMRequestProcessing(data=dict(request)) + with pytest.raises(ProxyRateLimitError) as exc_info: + await processor._pre_call_with_fallbacks( + request=self._chat_request(), general_settings={}, - version="test", - ) - if processor.data.get("model") == primary_model: - raise ProxyRateLimitError( - detail="TPM limit exceeded for gpt-4", - headers={"retry-after": "30"}, - ) - return processor.data, MagicMock() - - mock_router = MagicMock() - mock_router.fallbacks = [{primary_model: [fallback_model]}] - - with patch.object( - processor, - "common_processing_pre_call_logic", - side_effect=real_add_litellm_data_pre_call, - ): - data, logging_obj = await processor._pre_call_with_fallbacks( - request=request_mock, - general_settings={}, - proxy_logging_obj=MagicMock(), - user_api_key_dict=user_api_key_dict, + proxy_logging_obj=rig[0], + user_api_key_dict=key, version=None, - proxy_config=MagicMock(), + proxy_config=rig[2], user_model=None, user_temperature=None, user_request_timeout=None, user_max_tokens=None, user_api_base=None, - model=primary_model, + model=None, route_type="acompletion", - llm_router=mock_router, + llm_router=rig[1], ) - assert processor.data["model"] == fallback_model + assert rig[3] == [primary_model, primary_model, fallback_model] + assert exc_info.value.status_code == 429 + assert "Rate limit exceeded" in str(exc_info.value.detail) + assert exc_info.value.headers["retry-after"] + assert processor.data["model"] == primary_model + assert processor.data["litellm_logging_obj"].model == primary_model + assert processor.data["litellm_call_id"] @pytest.mark.asyncio - async def test_no_fallbacks_skips_snapshot(self): - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + async def test_fallback_lookup_uses_alias_resolved_model_group(self, monkeypatch: pytest.MonkeyPatch): + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + monkeypatch.setattr(litellm, "model_alias_map", {"my-alias": primary_model}) + key = self._otel_key(metadata={"model_rpm_limit": {primary_model: 1}}) + rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request = {"model": "my-alias", "messages": [{"role": "user", "content": "hi"}]} - processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + await self._pre_call(dict(request), key, rig) + _, (data, _) = await self._pre_call(dict(request), key, rig) - async def mock_pre_call_logic(**kwargs): - raise ProxyRateLimitError( - detail="TPM limit exceeded", - headers={"retry-after": "30"}, - ) - - mock_router = MagicMock() - mock_router.fallbacks = None - - with patch( # test-quality-ok: spying the snapshot seam is the only observable check that the no-fallback path skips it - "litellm.proxy.common_request_processing.independent_snapshot" - ) as snapshot_mock: - with patch.object( - processor, - "common_processing_pre_call_logic", - side_effect=mock_pre_call_logic, - ): - with pytest.raises(ProxyRateLimitError): - await processor._pre_call_with_fallbacks( - request=MagicMock(), - general_settings={}, - proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), - version=None, - proxy_config=MagicMock(), - user_model=None, - user_temperature=None, - user_request_timeout=None, - user_max_tokens=None, - user_api_base=None, - model="gpt-4", - route_type="acompletion", - llm_router=mock_router, - ) - - snapshot_mock.assert_not_called() + assert data["model"] == fallback_model + assert rig[3] == [primary_model, primary_model, fallback_model] class _RecordingSuccessLogger(CustomLogger): From 9faaf7f4d436b558dbc4d1ab25a5801269629d6d Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:15:13 +0000 Subject: [PATCH 063/187] refactor(proxy): assign the fallback model on the fresh snapshot instead of building a dict literal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7c9a296965c..d18fab1c9f0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2113,7 +2113,8 @@ class ProxyBaseLLMRequestProcessing: for fallback_model in fallback_models: if fallback_model == original_model: continue - self.data = {**independent_snapshot(pristine), "model": fallback_model} + self.data = independent_snapshot(pristine) + self.data["model"] = fallback_model try: return await self.common_processing_pre_call_logic( request=request, From 05ededf8a05c3b290c38e73a747bd602a49033d3 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:24:52 +0000 Subject: [PATCH 064/187] fix(otel): parent passthrough trace propagation on the legacy request span Pass user_api_key_dict.parent_otel_span into the outgoing W3C injection so the legacy otel callback propagates its litellm_request span, falling back to the otel_v2 request root span and then the ambient span. Extend the mapped unit tests to assert the propagated trace and span ids over real captured headers for HTTP and WebSocket passthrough with forwarding on and off. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 29 +++++------- .../pass_through_endpoints.py | 13 +++--- .../otel/test_otel_v2_components.py | 32 +++++++++---- .../test_pass_through_endpoints.py | 46 +++++++++++++------ 4 files changed, 74 insertions(+), 46 deletions(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 851097a17e7..3dabd34c81c 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -310,7 +310,10 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None: return _PROPAGATOR.extract(carrier) -def _outgoing_trace_context(inbound_headers: Mapping[str, str] | None = None) -> Context | None: +def _outgoing_trace_context(parent_span: object) -> Context | None: + if isinstance(parent_span, Span) and is_recordable_span(parent_span): + return context_from_span(parent_span) + root: Final = request_root_span() if root is not None: return context_from_span(root) @@ -318,27 +321,19 @@ def _outgoing_trace_context(inbound_headers: Mapping[str, str] | None = None) -> current: Final = get_current() if is_recordable_span(get_current_span(current)): return current - - if inbound_headers is None: - return None - inbound_context: Final = extract_traceparent(inbound_headers) - if inbound_context is None or not is_recordable_span(get_current_span(inbound_context)): - return None - return inbound_context + return None -def inject_trace_context( - headers: Mapping[str, str], - inbound_headers: Mapping[str, str] | None = None, -) -> dict[str, str]: +def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]: """``headers`` plus W3C ``traceparent``/``tracestate`` for the current request's span. - Parent preference: the anchored request root span, then the ambient active span, - then the trace context the caller sent inbound. Only trace context is injected, - never Baggage, so per-request identity baggage cannot leak upstream. Unchanged - when no valid span context exists anywhere. + Parent preference: the request span auth stashed on the key (the legacy + ``litellm_request`` SERVER span, or the FastAPI server span under otel_v2), then + the anchored request root span, then the ambient active span. Only trace context + is injected, never Baggage, so per-request identity baggage cannot leak upstream. + Unchanged when no valid span context exists anywhere. """ - context: Final = _outgoing_trace_context(inbound_headers) + context: Final = _outgoing_trace_context(parent_span) if context is None: return dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier carrier: Final = dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 6d44b4fa6dc..c8c2db0dbd6 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -985,7 +985,7 @@ async def pass_through_request( headers=headers, forward_headers=forward_headers, ) - headers = _with_trace_context(headers, inbound_headers=_safe_get_request_headers(request)) + headers = _with_trace_context(headers, parent_span=user_api_key_dict.parent_otel_span) requested_query_params: dict | None = query_params or dict(request.query_params) @@ -2161,12 +2161,12 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None: _WEBSOCKET_FORWARDED_HEADERS: Final = frozenset(("authorization", "x-api-key", "x-goog-user-project")) -def _with_trace_context(headers: Mapping[str, str], inbound_headers: Mapping[str, str]) -> dict[str, str]: +def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict[str, str]: try: from litellm.integrations.otel.plumbing.context import inject_trace_context except ImportError: return dict(headers) # mutable-ok: matches inject_trace_context's carrier return type - return inject_trace_context(headers, inbound_headers=inbound_headers) + return inject_trace_context(headers, parent_span=parent_span) async def websocket_passthrough_request( @@ -2211,16 +2211,15 @@ async def websocket_passthrough_request( await websocket.accept() verbose_proxy_logger.debug("WebSocket passthrough (%s): WebSocket connection accepted", endpoint) - incoming_headers: Final = dict(websocket.headers) # mutable-ok: propagator carrier - forwarded_headers: Final = { # mutable-ok: propagator carrier + forwarded_headers: Final = { # mutable-ok: one-shot upstream header dict, read as a Mapping **custom_headers, **{ header_name: header_value - for header_name, header_value in incoming_headers.items() + for header_name, header_value in websocket.headers.items() if forward_headers and header_name.lower() in _WEBSOCKET_FORWARDED_HEADERS }, } - upstream_headers: Final = _with_trace_context(forwarded_headers, inbound_headers=incoming_headers) + upstream_headers: Final = _with_trace_context(forwarded_headers, parent_span=user_api_key_dict.parent_otel_span) # Initialize logging object similar to HTTP passthrough team_callbacks: Final = _resolve_team_callback_wiring( diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index b1e65da1661..43546241d06 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -507,16 +507,32 @@ def test_inject_trace_context_uses_ambient_span_without_request_root(): assert propagated.get_span_context().span_id == ambient.get_span_context().span_id -def test_inject_trace_context_forwards_valid_inbound_context_without_span(): - inbound = {"traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"} - +def test_inject_trace_context_prefers_explicit_parent_span_over_root_and_ambient(): def run(): - result = ctx_mod.inject_trace_context({}, inbound_headers=inbound) - return get_current_span(TraceContextTextMapPropagator().extract(result)) + tracer = _test_tracer() + parent = tracer.start_span("litellm_request") + with tracer.start_as_current_span("ambient") as ambient: + ctx_mod.set_request_root_span(ambient) + result = ctx_mod.inject_trace_context({}, parent_span=parent) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return parent, ambient, propagated - propagated = ContextVarContext().run(run) - assert propagated.get_span_context().trace_id == int("0af7651916cd43dd8448eb211c80319c", 16) - assert propagated.get_span_context().span_id == int("b7ad6b7169203331", 16) + parent, ambient, propagated = ContextVarContext().run(run) + assert propagated.get_span_context().trace_id == parent.get_span_context().trace_id + assert propagated.get_span_context().span_id == parent.get_span_context().span_id + assert propagated.get_span_context().span_id != ambient.get_span_context().span_id + + +def test_inject_trace_context_skips_unusable_parent_span(): + def run(): + tracer = _test_tracer() + with tracer.start_as_current_span("ambient") as ambient: + result = ctx_mod.inject_trace_context({}, parent_span=object()) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return ambient, propagated + + ambient, propagated = ContextVarContext().run(run) + assert propagated.get_span_context().span_id == ambient.get_span_context().span_id def test_inject_trace_context_returns_headers_unchanged_without_context(): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 0603b5cb2ff..d6b9a782f99 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -51,7 +51,7 @@ MESSAGE_START_SSE_FRAME = b'event: message_start\ndata: {"type": "message_start" def test_with_trace_context_without_opentelemetry(monkeypatch: pytest.MonkeyPatch): monkeypatch.setitem(sys.modules, "litellm.integrations.otel.plumbing.context", None) - headers = _with_trace_context({"authorization": "x"}, {}) + headers = _with_trace_context({"authorization": "x"}, parent_span=None) assert headers == {"authorization": "x"} assert "traceparent" not in headers @@ -4282,11 +4282,11 @@ def _relay_client_request(method="GET"): @pytest.mark.asyncio -async def test_pass_through_request_propagates_active_trace_context(): +@pytest.mark.parametrize("span_source", ["auth_parent_span", "ambient_span"]) +async def test_pass_through_request_propagates_active_trace_context(span_source: str): from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import get_current_span from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator - from litellm.proxy._types import UserAPIKeyAuth captured: dict[str, httpx.Headers] = {} @@ -4295,17 +4295,23 @@ async def test_pass_through_request_propagates_active_trace_context(): return httpx.Response(200, json={"ok": True}, request=upstream_request) fake_client, cleanup = _inject_fake_passthrough_client(httpx.MockTransport(transport_handler), timeout=None) + tracer = TracerProvider().get_tracer("test") try: with ExitStack() as stack: _enter_relay_logging_mocks(stack, {}) - tracer = TracerProvider().get_tracer("test") - with tracer.start_as_current_span("passthrough") as span: - response = await pass_through_request( - request=_relay_client_request(method="POST"), - target="http://internal-api.test/v1/generate", - custom_headers={}, - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), - ) + if span_source == "auth_parent_span": + span = tracer.start_span("litellm_request") + stack.callback(span.end) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", parent_otel_span=span) + else: + span = stack.enter_context(tracer.start_as_current_span("passthrough")) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") + response = await pass_through_request( + request=_relay_client_request(method="POST"), + target="http://internal-api.test/v1/generate", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + ) finally: cleanup() await fake_client.aclose() @@ -4313,6 +4319,7 @@ async def test_pass_through_request_propagates_active_trace_context(): assert response.status_code == 200 propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) assert propagated.get_span_context().trace_id == span.get_span_context().trace_id + assert propagated.get_span_context().span_id == span.get_span_context().span_id @pytest.mark.asyncio @@ -4913,7 +4920,10 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): @pytest.mark.asyncio @pytest.mark.parametrize("forward_headers", [True, False]) -async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch, forward_headers: bool): +@pytest.mark.parametrize("span_source", ["auth_parent_span", "ambient_span"]) +async def test_websocket_passthrough_propagates_active_trace_context( + monkeypatch, forward_headers: bool, span_source: str +): from opentelemetry.sdk.trace import TracerProvider from opentelemetry.trace import get_current_span from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator @@ -4954,12 +4964,19 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch "litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER", mock_worker, ) - with tracer.start_as_current_span("websocket_passthrough") as span: + with ExitStack() as stack: + if span_source == "auth_parent_span": + span = tracer.start_span("litellm_request") + stack.callback(span.end) + user_api_key_dict = UserAPIKeyAuth(parent_otel_span=span) + else: + span = stack.enter_context(tracer.start_as_current_span("websocket_passthrough")) + user_api_key_dict = UserAPIKeyAuth() await websocket_passthrough_request( websocket=websocket, target="wss://upstream.example.test/v1/realtime", custom_headers={}, - user_api_key_dict=UserAPIKeyAuth(), + user_api_key_dict=user_api_key_dict, forward_headers=forward_headers, endpoint="/realtime", accept_websocket=True, @@ -4967,6 +4984,7 @@ async def test_websocket_passthrough_propagates_active_trace_context(monkeypatch propagated = get_current_span(TraceContextTextMapPropagator().extract(captured["headers"])) assert propagated.get_span_context().trace_id == span.get_span_context().trace_id + assert propagated.get_span_context().span_id == span.get_span_context().span_id assert captured["headers"].get("authorization") == ("Bearer client" if forward_headers else None) From 8899583d06ea4f570136cdcdecceb41f2de652f6 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:33:39 +0000 Subject: [PATCH 065/187] docs(otel): tighten inject_trace_context docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 3dabd34c81c..57e220815b5 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -325,13 +325,11 @@ def _outgoing_trace_context(parent_span: object) -> Context | None: def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]: - """``headers`` plus W3C ``traceparent``/``tracestate`` for the current request's span. + """``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span. - Parent preference: the request span auth stashed on the key (the legacy - ``litellm_request`` SERVER span, or the FastAPI server span under otel_v2), then - the anchored request root span, then the ambient active span. Only trace context - is injected, never Baggage, so per-request identity baggage cannot leak upstream. - Unchanged when no valid span context exists anywhere. + Parent preference: ``parent_span`` (the request span auth stashed on the key), then + the anchored request root span, then the ambient active span. Only trace context is + injected, never Baggage. Unchanged when no valid span exists anywhere. """ context: Final = _outgoing_trace_context(parent_span) if context is None: From 391da46e2cdcef46491519ba2814b5d38c752285 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 07:59:33 +0000 Subject: [PATCH 066/187] test: type the v3 limiter rig and otel key helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/test_common_request_processing.py | 34 ++++++++++++------- 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f689dd62df6..7d9a78fd981 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -44,6 +44,7 @@ from litellm.proxy.common_request_processing import ( create_response, ) from litellm.proxy.dd_span_tagger import DDSpanTagger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth from litellm.proxy.utils import ProxyLogging @@ -6383,15 +6384,17 @@ class TestPreCallWithFallbacksOnLocalRateLimit: limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) limiter_models: list[str] = [] - async def run_limiter(**kwargs): - limiter_models.append(kwargs["data"]["model"]) + async def run_limiter( + user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str + ) -> dict[str, object]: + limiter_models.append(str(data["model"])) await limiter.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=DualCache(), - data=kwargs["data"], - call_type=kwargs["call_type"], + data=data, + call_type=call_type, ) - return kwargs["data"] + return data proxy_logging_obj = MagicMock(spec=ProxyLogging) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=run_limiter) @@ -6406,11 +6409,18 @@ class TestPreCallWithFallbacksOnLocalRateLimit: return proxy_logging_obj, router, proxy_server.ProxyConfig(), limiter_models @staticmethod - def _otel_key(**limits) -> ProxyUserAPIKeyAuth: + def _otel_key( + rpm_limit: int | None = None, model_rpm_limit: dict[str, int] | None = None + ) -> ProxyUserAPIKeyAuth: from opentelemetry.sdk.trace import TracerProvider span = TracerProvider().get_tracer("test").start_span("proxy-request") - return ProxyUserAPIKeyAuth(api_key="hashed-key", parent_otel_span=span, **limits) + return ProxyUserAPIKeyAuth( + api_key="hashed-key", + parent_otel_span=span, + rpm_limit=rpm_limit, + metadata={"model_rpm_limit": model_rpm_limit} if model_rpm_limit else {}, + ) @staticmethod def _chat_request() -> Request: @@ -6418,10 +6428,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: async def _pre_call( self, - data: dict, + data: dict[str, object], user_api_key_dict: ProxyUserAPIKeyAuth, rig: tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]], - ) -> tuple[ProxyBaseLLMRequestProcessing, tuple[dict, object]]: + ) -> tuple[ProxyBaseLLMRequestProcessing, tuple[dict[str, object], LiteLLMLoggingObj]]: proxy_logging_obj, router, proxy_config, _ = rig processor = ProxyBaseLLMRequestProcessing(data=data) result = await processor._pre_call_with_fallbacks( @@ -6450,10 +6460,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: never deep-copies the span (the ``cannot pickle '_thread.RLock'`` 500).""" primary_model = "gpt-4.1" fallback_model = "gpt-4.1-mini" - key = self._otel_key(metadata={"model_rpm_limit": {primary_model: 1}}) + key = self._otel_key(model_rpm_limit={primary_model: 1}) rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) - def client_request() -> dict: + def client_request() -> dict[str, object]: return { "model": primary_model, "messages": [{"role": "user", "content": "hi"}], @@ -6520,7 +6530,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: primary_model = "gpt-4.1" fallback_model = "gpt-4.1-mini" monkeypatch.setattr(litellm, "model_alias_map", {"my-alias": primary_model}) - key = self._otel_key(metadata={"model_rpm_limit": {primary_model: 1}}) + key = self._otel_key(model_rpm_limit={primary_model: 1}) rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) request = {"model": "my-alias", "messages": [{"role": "user", "content": "hi"}]} From b82de95625cb807e892a223f90004d46c43589c0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 08:04:27 +0000 Subject: [PATCH 067/187] refactor(passthrough): bind trace-enriched headers to a Final local Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../pass_through_endpoints.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index c8c2db0dbd6..7d2fe8c45a9 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -985,7 +985,7 @@ async def pass_through_request( headers=headers, forward_headers=forward_headers, ) - headers = _with_trace_context(headers, parent_span=user_api_key_dict.parent_otel_span) + upstream_headers: Final = _with_trace_context(headers, parent_span=user_api_key_dict.parent_otel_span) requested_query_params: dict | None = query_params or dict(request.query_params) @@ -1019,7 +1019,7 @@ async def pass_through_request( verbose_proxy_logger.debug( "Pass through endpoint sending request to \nURL %s\nheaders: %s\nbody: %s\n", url, - headers, + upstream_headers, _parsed_body, ) @@ -1257,7 +1257,7 @@ async def pass_through_request( additional_args={ "complete_input_dict": _parsed_body, "api_base": str(logging_url), - "headers": headers, + "headers": upstream_headers, }, ) stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body( @@ -1274,7 +1274,7 @@ async def pass_through_request( request=request, async_client=async_client, url=url, - headers=headers, + headers=upstream_headers, requested_query_params=requested_query_params, stream=True, ) @@ -1286,7 +1286,7 @@ async def pass_through_request( request.method, url, params=requested_query_params, - headers=headers, + headers=upstream_headers, content=state_raw_body, ) if state_raw_body is not None @@ -1294,7 +1294,7 @@ async def pass_through_request( request.method, url, params=requested_query_params, - headers=headers, + headers=upstream_headers, json=_parsed_body, ) ) @@ -1371,7 +1371,7 @@ async def pass_through_request( raw_body_request: Final = async_client.build_request( request.method, url, - headers=headers, + headers=upstream_headers, params=requested_query_params, content=state_raw_body, ) @@ -1381,7 +1381,7 @@ async def pass_through_request( request=request, async_client=async_client, url=url, - headers=headers, + headers=upstream_headers, requested_query_params=requested_query_params, _parsed_body=_parsed_body, forward_multipart=is_multipart, From 9974cf4bf817331053649fccc1f1b9e00d61d570 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 08:17:58 +0000 Subject: [PATCH 068/187] fix(proxy): honor key-level disable_fallbacks after first pre-call pass Key metadata disable_fallbacks only lands on data during add_key_level_controls, so the local rate-limit fallback retry now rechecks it post pre-call. Also use a real UserAPIKeyAuth in the skip pre-call test since the path reads router_settings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 7 ++++- .../test_response_polling_pre_call_checks.py | 2 +- .../proxy/test_common_request_processing.py | 30 +++++++++++++++++-- 3 files changed, 35 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d18fab1c9f0..36b49ef064c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2093,7 +2093,12 @@ class ProxyBaseLLMRequestProcessing: except ProxyRateLimitError as original_exc: rate_limited_data: Final = self.data original_model: Final = rate_limited_data.get("model") - if pristine is None or not configured_fallbacks or not isinstance(original_model, str): + if ( + pristine is None + or not configured_fallbacks + or rate_limited_data.get("disable_fallbacks") + or not isinstance(original_model, str) + ): raise fallback_models: Final = self._resolve_fallback_models( diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py index 459834d0fd2..9f1a228855e 100644 --- a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py +++ b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py @@ -87,7 +87,7 @@ class TestSkipPreCallLogic: await processor.base_process_llm_request( request=MagicMock(spec=Request), fastapi_response=MagicMock(spec=Response), - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + user_api_key_dict=UserAPIKeyAuth(), route_type="aresponses", proxy_logging_obj=mock_proxy_logging, llm_router=MagicMock(), diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 7d9a78fd981..e9d74a46daa 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6410,7 +6410,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: @staticmethod def _otel_key( - rpm_limit: int | None = None, model_rpm_limit: dict[str, int] | None = None + rpm_limit: int | None = None, + model_rpm_limit: dict[str, int] | None = None, + disable_fallbacks: bool = False, ) -> ProxyUserAPIKeyAuth: from opentelemetry.sdk.trace import TracerProvider @@ -6419,7 +6421,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: api_key="hashed-key", parent_otel_span=span, rpm_limit=rpm_limit, - metadata={"model_rpm_limit": model_rpm_limit} if model_rpm_limit else {}, + metadata={ + **({"model_rpm_limit": model_rpm_limit} if model_rpm_limit else {}), + **({"disable_fallbacks": True} if disable_fallbacks else {}), + }, ) @staticmethod @@ -6540,6 +6545,27 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert data["model"] == fallback_model assert rig[3] == [primary_model, primary_model, fallback_model] + @pytest.mark.asyncio + async def test_key_metadata_disable_fallbacks_returns_429_instead_of_retrying( + self, monkeypatch: pytest.MonkeyPatch + ): + """``disable_fallbacks`` set in key metadata only lands on ``data`` during the first + pre-call pass (``add_key_level_controls``), so it must be honored after that pass.""" + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + primary_model = "gpt-4.1" + fallback_model = "gpt-4.1-mini" + key = self._otel_key(model_rpm_limit={primary_model: 1}, disable_fallbacks=True) + rig = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}]) + request = {"model": primary_model, "messages": [{"role": "user", "content": "hi"}]} + + await self._pre_call(dict(request), key, rig) + with pytest.raises(ProxyRateLimitError) as exc_info: + await self._pre_call(dict(request), key, rig) + + assert exc_info.value.status_code == 429 + assert rig[3] == [primary_model, primary_model] + class _RecordingSuccessLogger(CustomLogger): def __init__(self): From cd594f104af2ab72242f6512e3b6aa240df658a4 Mon Sep 17 00:00:00 2001 From: shivam Date: Wed, 16 Sep 2026 08:34:53 +0000 Subject: [PATCH 069/187] fix(otel): drop stale trace headers before injecting passthrough trace context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/plumbing/context.py | 5 ++++- .../otel/test_otel_v2_components.py | 20 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 57e220815b5..1282e654365 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -26,6 +26,7 @@ if TYPE_CHECKING: from litellm.integrations.otel.model.destination import OtelDestination _PROPAGATOR: Final = TraceContextTextMapPropagator() +_W3C_TRACE_HEADERS: Final = frozenset(("traceparent", "tracestate")) # The request's root span — the FastAPI-owned SERVER span — captured ONCE when the # proxy first resolves it, so request-level spans (the LLM call, guardrails) can @@ -334,7 +335,9 @@ def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) context: Final = _outgoing_trace_context(parent_span) if context is None: return dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier - carrier: Final = dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier + carrier: Final = { # mutable-ok: OpenTelemetry propagator requires a mutable carrier + key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS + } _PROPAGATOR.inject(carrier, context=context) return carrier diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index 43546241d06..72f6213e880 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -507,6 +507,26 @@ def test_inject_trace_context_uses_ambient_span_without_request_root(): assert propagated.get_span_context().span_id == ambient.get_span_context().span_id +def test_inject_trace_context_replaces_stale_trace_headers(): + def run(): + tracer = _test_tracer() + with tracer.start_as_current_span("ambient") as ambient: + headers = { + "Traceparent": "00-" + "a" * 32 + "-" + "b" * 16 + "-01", + "Tracestate": "vendor=old", + "x-keep": "1", + } + result = ctx_mod.inject_trace_context(headers) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return result, ambient, propagated + + result, ambient, propagated = ContextVarContext().run(run) + assert sum(key.lower() == "traceparent" for key in result) == 1 + assert not any(key.lower() == "tracestate" for key in result) + assert result["x-keep"] == "1" + assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id + + def test_inject_trace_context_prefers_explicit_parent_span_over_root_and_ambient(): def run(): tracer = _test_tracer() From 2d40254b57a62be9bce95586cb902b90c8505c4c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 02:15:46 -0700 Subject: [PATCH 070/187] feat(e2e): key the provider cache per test and mount Bedrock behind it The exact-request cache reused 5% of routed traffic (build 218: 19 hits, 350 misses) because every test salts its prompt with a fresh unique_marker(), so the same test could never match itself across builds. It also routed only openai and anthropic, while the week's flakiness was Bedrock. Key is now HMAC(test id + method + URL + headers + body, with every unique_marker() token replaced by a placeholder, + FIFO slot index). The slot index is what keeps two marker-only-different calls in one test on two recordings and therefore two provider response ids, so spend rows still reconcile one per invocation. A call outside any test is not cacheable. Bedrock gets a region-qualified mount and SigV4 re-signing, since the edge rewrites the Host the proxy signed. Signature headers are excluded from the key for signing mounts only, because x-amz-date would otherwise make every Bedrock request a permanent miss; every other mount still keys on its credentials whole. Only Anthropic-on-Bedrock chat deployments route: embeddings, image generation, rerank and realtime keep their direct path, and so do deployments carrying their own aws_role_name or static keys, whose whole point is to prove the product's assume-role chain rather than the runner's. The two eventstream actions bypass the cache and go live, still signed. Counters are now attributed per mount as well as in total, so a build can report a per-provider hit rate instead of one number. --- .../test_provider_cache.py | 509 +++++++++++++++--- tests/e2e/fixture_canonical.py | 5 +- tests/e2e/models.py | 1 + tests/e2e/provider_cache.py | 171 ++++-- tests/e2e/provider_cache_routing.py | 47 +- tests/e2e/provider_edge.py | 62 ++- tests/e2e/provider_edge_bedrock.py | 72 +++ tests/e2e/test_provider_edge.py | 19 +- 8 files changed, 768 insertions(+), 118 deletions(-) create mode 100644 tests/e2e/provider_edge_bedrock.py diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 828227ed239..5d35981344d 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -7,7 +7,7 @@ import subprocess import threading import time import uuid -from collections.abc import Generator +from collections.abc import Generator, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, replace @@ -18,21 +18,58 @@ from urllib.parse import urlsplit import pytest from e2e_http import NetworkError, PreparedForward, RawResponse, StreamChunk, StreamHead, forward, prepare_forward -from models import LiteLLMParamsBody -from provider_cache import CacheEdge, CacheHit, CaptureLease, exact_key, successful_response +from models import LiteLLMParamsBody, ModelMode +from botocore.credentials import Credentials +from provider_cache import ( + SIGNATURE_HEADERS, + CacheEdge, + CacheHit, + CaptureLease, + ResponseStore, + cacheable_endpoint, + request_identity, + slotted_key, + successful_response, +) from provider_cache_redis import PUBLISH, RedisCommands, RedisResponseStore, configured_cache, redis_store from provider_cache_routing import LIVE_PROVIDER_REQUIRED, route_cache_model -from provider_edge import configured_cache_backend, start_provider_edge +from fixture_mode import SESSION_TEST_KEY +from provider_edge import EDGE_MOUNTS, configured_cache_backend, resolve_mount, start_provider_edge +from provider_edge_bedrock import bedrock_signer from redis.exceptions import ConnectionError as RedisConnectionError SECRET: Final = b"synthetic-cache-hmac-key-for-tests" BODY: Final = b'{"model":"test","messages":[{"role":"user","content":"hello"}]}' SUCCESS: Final = b'{"id":"provider-fixed-id","choices":[{"message":{"content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}' HEADERS: Final = {"content-type": "application/json", "authorization": "Bearer synthetic-account-one"} +TEST_KEY: Final = "tests/e2e/synthetic_suite.py::TestCase::test_case" +OTHER_TEST_KEY: Final = "tests/e2e/synthetic_suite.py::TestCase::test_other_case" + + +def marked(marker: str) -> bytes: + """One request body shaped like the suite's own: a fixed prompt salted with a + 12-lowercase-hex ``unique_marker()`` token, fresh on every run.""" + return b'{"model":"test","messages":[{"role":"user","content":"hello %s"}]}' % marker.encode() + + +MARKED: Final = marked("0a1b2c3d4e5f") +BEDROCK_MOUNT: Final = "bedrock/us-east-1" +BEDROCK_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1%3A0" +BEDROCK_BODY: Final = b'{"messages":[{"role":"user","content":[{"text":"hello 0a1b2c3d4e5f"}]}]}' +CONVERSE_SUCCESS: Final = ( + b'{"output":{"message":{"role":"assistant","content":[{"text":"hi"}]}},' + b'"stopReason":"end_turn","usage":{"inputTokens":1,"outputTokens":1,"totalTokens":2}}' +) +INVOKE_SUCCESS: Final = ( + b'{"id":"msg_synthetic","type":"message","role":"assistant",' + b'"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}' +) +STATIC_CREDENTIALS: Final = Credentials("AKIAIOSFODNN7EXAMPLE", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY") class Provider(ThreadingHTTPServer): hits: tuple[tuple[str, bytes], ...] = () + authorizations: tuple[str, ...] = () response: bytes = SUCCESS status: int = 200 delay: float = 0 @@ -49,6 +86,7 @@ class Handler(BaseHTTPRequestHandler): assert isinstance(server, Provider) body: Final = self.rfile.read(int(self.headers.get("content-length", "0"))) server.hits += ((self.path, body),) + server.authorizations += (self.headers.get("authorization", ""),) time.sleep(server.delay) self.send_response(server.status) if server.stream: @@ -122,6 +160,29 @@ def store(redis_url: str) -> RedisResponseStore: return redis_store(redis_url, "test-" + uuid.uuid4().hex) +def cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: + """A cache edge standing in for one pytest process. A fresh instance over the + same store is the next build running the same test: the recordings survive, + the per-test FIFO slot counters start over.""" + return CacheEdge(store, SECRET, test_key=lambda: test_key) + + +def slot_key( + url: str, slot: int = 0, body: bytes | None = BODY, + headers: dict[str, str] = HEADERS, test_key: str = TEST_KEY, +) -> str: + prepared: Final = prepare_forward("POST", url, headers, body) + assert isinstance(prepared, PreparedForward) + return slotted_key(SECRET, request_identity(SECRET, test_key, "POST", url, prepared.headers, body), slot) + + +def bedrock_cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: + return CacheEdge( + store, SECRET, test_key=lambda: test_key, + signers={BEDROCK_MOUNT: bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS)}, + ) + + @contextmanager def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]: upstream: Final = f"http://127.0.0.1:{provider.server_port}" @@ -132,34 +193,54 @@ def edge(cache: CacheEdge, provider: Provider) -> Generator[str, None, None]: running.shutdown() +@contextmanager +def bedrock_edge(cache: CacheEdge, provider: Provider, action: str = "converse") -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + running: Final = start_provider_edge(cache, mounts={BEDROCK_MOUNT: upstream}) + try: + yield f"{running.edge.api_base(BEDROCK_MOUNT)}/model/{BEDROCK_MODEL}/{action}" + finally: + running.shutdown() + + def call(url: str, body: bytes = BODY, headers: dict[str, str] = HEADERS) -> RawResponse: result: Final = forward("POST", url, headers=headers, body=body, timeout=5) assert isinstance(result, RawResponse), result return result -def test_success_is_reusable_across_fresh_edges(store: RedisResponseStore, provider: Provider) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: +def test_repeated_call_takes_its_own_slot_and_both_replay_next_run( + store: RedisResponseStore, provider: Provider, +) -> None: + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS assert call(url).body == SUCCESS - with edge(CacheEdge(store, SECRET), provider) as other: + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as other: assert call(other).body == SUCCESS - assert len(provider.hits) == 1 + assert call(other).body == SUCCESS + assert len(provider.hits) == 2 @pytest.mark.parametrize("body", [BODY + b" ", BODY.replace(b"hello", b"Hello"), BODY.replace(b"test", b"test2")]) def test_any_body_change_calls_live(store: RedisResponseStore, provider: Provider, body: bytes) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: call(url, body) + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as url: call(url, body) assert len(provider.hits) == 2 @pytest.mark.parametrize("name,value", [("authorization", "Bearer another-account"), ("x-request-id", "one"), ("anthropic-version", "new")]) def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provider, name: str, value: str) -> None: - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: call(url, headers=HEADERS | {name: value}) call(url + "?x=1") assert len(provider.hits) == 3 @@ -169,36 +250,59 @@ def test_changed_header_cannot_reuse(store: RedisResponseStore, provider: Provid def test_failed_provider_responses_never_enter_cache(store: RedisResponseStore, provider: Provider, status: int, response: bytes) -> None: provider.status = status provider.response = response - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: assert call(url).status_code == status + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: assert call(url).body == response assert len(provider.hits) == 2 def test_cookie_setting_success_is_reused_without_the_cookie(store: RedisResponseStore, provider: Provider) -> None: provider.cookie = "__cf_bm=synthetic-bot-management; Path=/; HttpOnly; Secure" - with edge(CacheEdge(store, SECRET), provider) as url: - replies: Final = tuple(call(url) for _ in range(2)) + with edge(cache_edge(store), provider) as url: + live: Final = call(url) + with edge(cache_edge(store), provider) as url: + replayed: Final = call(url) assert len(provider.hits) == 1 - assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in replies) + assert all(reply.body == SUCCESS and "set-cookie" not in reply.headers for reply in (live, replayed)) def test_expiry_does_not_slide(store: RedisResponseStore, provider: Provider) -> None: short: Final = replace(store, lifetime_ms=250) - with edge(CacheEdge(short, SECRET), provider) as url: - call(url) - call(url) - time.sleep(0.3) - call(url) - call(url) + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + + def drain() -> None: + head = cache_edge(short).forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS + + drain() + assert len(provider.hits) == 1 + drain() + assert len(provider.hits) == 1 + time.sleep(0.3) + drain() assert len(provider.hits) == 2 -def test_concurrent_requests_publish_atomically(store: RedisResponseStore, provider: Provider) -> None: +def test_concurrent_builds_publish_one_recording_atomically( + store: RedisResponseStore, provider: Provider, +) -> None: + """Five processes running the same test at the same time all reach slot 0 of + one key, which is the only way the capture lease is contended now that a + repeat inside a single test takes its own slot.""" provider.delay = 0.15 - with edge(CacheEdge(store, SECRET), provider) as url: - with ThreadPoolExecutor(max_workers=5) as executor: - replies: Final = tuple(executor.map(lambda _: call(url).body, range(5))) + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + edges: Final = tuple(cache_edge(store) for _ in range(5)) + + def drain(cache: CacheEdge) -> bytes: + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + return b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) + + with ThreadPoolExecutor(max_workers=5) as executor: + replies: Final = tuple(executor.map(drain, edges)) assert replies == (SUCCESS,) * 5 assert len(provider.hits) == 1 @@ -231,9 +335,9 @@ def test_stream_completion_controls_publication(store: RedisResponseStore, provi provider.stream = True provider.truncated = truncated provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n' - with edge(CacheEdge(store, SECRET), provider) as url: - for _ in range(2): - result: Final = forward("POST", url, headers=HEADERS, body=BODY, timeout=5) + for _ in range(2): + with edge(cache_edge(store), provider) as url: + result = forward("POST", url, headers=HEADERS, body=BODY, timeout=5) if truncated: assert isinstance(result, NetworkError) else: @@ -246,9 +350,9 @@ def test_store_outage_preserves_provider_success(provider: Provider) -> None: probe.bind(("127.0.0.1", 0)) port: Final = probe.getsockname()[1] unavailable: Final = redis_store(f"redis://127.0.0.1:{port}/0", "unavailable") - with edge(CacheEdge(unavailable, SECRET), provider) as url: - assert call(url).body == SUCCESS - assert call(url).body == SUCCESS + for _ in range(2): + with edge(cache_edge(unavailable), provider) as url: + assert call(url).body == SUCCESS assert len(provider.hits) == 2 @@ -267,7 +371,8 @@ def test_old_lease_cannot_overwrite_new_owner(store: RedisResponseStore) -> None def test_identity_preserves_values_and_never_contains_credentials() -> None: variants: Final = (b'{}', b'{"a":null}', b'{"a":false}', b'{"a":0}', b'{"a":0.0}', b'{"a":"0"}', b' { }', None, b'') - keys: Final = tuple(exact_key(SECRET, "POST", "https://example.invalid/v1/chat/completions", HEADERS, body) for body in variants) + url: Final = "https://example.invalid/v1/chat/completions" + keys: Final = tuple(request_identity(SECRET, TEST_KEY, "POST", url, HEADERS, body) for body in variants) assert len(set(keys)) == len(variants) assert all(len(key) == 64 and "synthetic-account" not in key for key in keys) @@ -275,21 +380,22 @@ def test_identity_preserves_values_and_never_contains_credentials() -> None: @pytest.mark.parametrize("payload", [b"corrupt response", '{"response":"{}","signature":"é"}'.encode()]) def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisResponseStore, provider: Provider, payload: bytes) -> None: upstream: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - prepared: Final = prepare_forward("POST", upstream, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", upstream, prepared.headers, BODY) + key: Final = slot_key(upstream) lease: Final = store.lookup(key) assert isinstance(lease, CaptureLease) assert store.publish(key, lease, payload) - cache: Final = CacheEdge(store, SECRET) - for _ in range(2): - head = cache.forward("POST", upstream, HEADERS, BODY, 5) + caches: Final = tuple(cache_edge(store) for _ in range(2)) + for cache in caches: + head = cache.forward("openai", "POST", upstream, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS assert len(provider.hits) == 1 - assert dict(cache.counters.counts) == { - "corrupt": 1, "misses": 1, "upstream_attempts": 1, "writes": 1, "hits": 1, + assert dict(caches[0].counters.counts) == { + "corrupt": 1, "mount:openai:corrupt": 1, "misses": 1, "mount:openai:misses": 1, + "upstream_attempts": 1, "mount:openai:upstream_attempts": 1, + "writes": 1, "mount:openai:writes": 1, } + assert dict(caches[1].counters.counts) == {"hits": 1, "mount:openai:hits": 1} @pytest.mark.parametrize("payload", [ @@ -301,22 +407,242 @@ def test_corrupt_entry_is_replaced_by_same_successful_request(store: RedisRespon def test_malformed_success_stream_is_never_cached(store: RedisResponseStore, provider: Provider, payload: bytes) -> None: provider.stream = True provider.response = payload - with edge(CacheEdge(store, SECRET), provider) as url: - assert call(url).body == payload - assert call(url).body == payload + for _ in range(2): + with edge(cache_edge(store), provider) as url: + assert call(url).body == payload assert len(provider.hits) == 2 +def test_requests_differing_only_by_marker_share_one_recording_per_slot( + store: RedisResponseStore, provider: Provider, +) -> None: + """The whole point of the canonical key. Every e2e test salts its prompt with + a fresh ``unique_marker()``, so before this the same test could never reuse + anything across builds. The second run mints markers it has never sent, which + is what a later build actually does, and must still serve both from the two + slots the first run recorded.""" + with edge(cache_edge(store), provider) as url: + assert call(url, MARKED).body == SUCCESS + assert call(url, marked("f5e4d3c2b1a0")).body == SUCCESS + assert len(provider.hits) == 2 + with edge(cache_edge(store), provider) as url: + assert call(url, marked("7c6b5a493827")).body == SUCCESS + assert call(url, marked("1122334455ff")).body == SUCCESS + assert len(provider.hits) == 2 + + +@pytest.mark.parametrize("body", [ + b'{"model":"test","messages":[{"role":"user","content":"hello 0a1b2c3d4e5"}]}', + b'{"model":"test","messages":[{"role":"user","content":"hello 0a1b2c3d4e5f0"}]}', + b'{"model":"test","messages":[{"role":"user","content":"hello 0A1B2C3D4E5F"}]}', + b'{"model":"0a1b2c3d4e5f","messages":[{"role":"user","content":"hello"}]}', +]) +def test_a_token_that_is_not_a_marker_keeps_its_own_key( + store: RedisResponseStore, provider: Provider, body: bytes, +) -> None: + """Too short, too long, upper case, or in another field: none of these is the + 12-lowercase-hex token ``unique_marker`` mints, so none may fold onto it.""" + with edge(cache_edge(store), provider) as url: + call(url, MARKED) + assert len(provider.hits) == 1 + with edge(cache_edge(store), provider) as url: + call(url, body) + assert len(provider.hits) == 2 + + +def test_another_test_never_reuses_this_tests_recording( + store: RedisResponseStore, provider: Provider, +) -> None: + with edge(cache_edge(store), provider) as url: + call(url) + assert len(provider.hits) == 1 + with edge(cache_edge(store, OTHER_TEST_KEY), provider) as url: + call(url) + assert len(provider.hits) == 2 + with edge(cache_edge(store, OTHER_TEST_KEY), provider) as url: + call(url) + assert len(provider.hits) == 2 + + +def test_calls_outside_any_test_are_never_cached( + store: RedisResponseStore, provider: Provider, +) -> None: + url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" + cache: Final = CacheEdge(store, SECRET, test_key=lambda: SESSION_TEST_KEY) + for _ in range(2): + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) + assert isinstance(head, StreamHead) + assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS + assert len(provider.hits) == 2 + assert dict(cache.counters.counts) == { + "bypass": 2, "mount:openai:bypass": 2, + "upstream_attempts": 2, "mount:openai:upstream_attempts": 2, + } + + +def test_counters_attribute_every_outcome_to_its_mount( + store: RedisResponseStore, provider: Provider, +) -> None: + """The build report needs per-provider hit counts, and the flat totals cannot + supply them. Anthropic is served a chat-shaped body here, which its validator + rejects, so one mount writes and the other does not.""" + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + cache: Final = cache_edge(store) + running: Final = start_provider_edge(cache, mounts={"openai": upstream, "anthropic": upstream}) + try: + call(running.edge.api_base("openai") + "/v1/chat/completions") + call(running.edge.api_base("anthropic") + "/v1/messages") + finally: + running.shutdown() + counts: Final = dict(cache.counters.counts) + assert counts["misses"] == 2 + assert counts["mount:openai:misses"] == 1 and counts["mount:anthropic:misses"] == 1 + assert counts["mount:openai:writes"] == 1 and "mount:anthropic:writes" not in counts + assert counts["mount:anthropic:rejected"] == 1 and "mount:openai:rejected" not in counts + + +class TestBedrockSigning: + """Bedrock is the reason the edge could not mount it before: SigV4 covers the + Host header, so forwarding through a rewritten api_base invalidates the + proxy's signature. The edge mints its own over the upstream URL instead.""" + + def test_the_proxys_signature_is_replaced_not_forwarded(self) -> None: + signer: Final = bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS) + signed: Final = signer( + "POST", + f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/converse", + {"content-type": "application/json", "Authorization": "AWS4-HMAC-SHA256 Credential=PROXY/...", + "X-Amz-Date": "19700101T000000Z", "X-Amz-Security-Token": "proxy-session-token"}, + BEDROCK_BODY, + ) + assert "PROXY" not in str(signed) and "proxy-session-token" not in str(signed) + assert signed["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") + assert "/us-east-1/bedrock/aws4_request" in signed["Authorization"] + assert signed["X-Amz-Date"] != "19700101T000000Z" + assert signed["content-type"] == "application/json" + + def test_the_signed_url_reaches_the_wire_byte_for_byte(self) -> None: + """SigV4 hashes the canonical URI, so if the HTTP layer re-encoded the + colon in an inference-profile id after signing, every call would fail + with a signature mismatch rather than anything that names the cause.""" + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/converse" + signer: Final = bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS) + prepared: Final = prepare_forward("POST", url, signer("POST", url, dict(HEADERS), BEDROCK_BODY), BEDROCK_BODY) + assert isinstance(prepared, PreparedForward) + assert urlsplit(prepared.url).path == urlsplit(url).path + + def test_signature_headers_are_excluded_from_the_key( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """A real signature is fresh on every call, so keying on it would make + every Bedrock request a permanent miss. The stub signer here varies its + stamp per call on purpose: the real one only varies once a second, which + would let this pass by luck when it should fail.""" + provider.response = CONVERSE_SUCCESS + stamps: Final = iter(("20260101T000000Z", "20260102T111111Z")) + + def varying(method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> dict[str, str]: + return dict(headers) | {"authorization": f"AWS4-HMAC-SHA256 {url}", "x-amz-date": next(stamps)} + + def signing_edge() -> CacheEdge: + return CacheEdge(store, SECRET, test_key=lambda: TEST_KEY, signers={BEDROCK_MOUNT: varying}) + + for _ in range(2): + with bedrock_edge(signing_edge(), provider) as url: + assert call(url, BEDROCK_BODY).body == CONVERSE_SUCCESS + assert len(provider.hits) == 1 + assert provider.authorizations[0] == ( + f"AWS4-HMAC-SHA256 http://127.0.0.1:{provider.server_port}/model/{BEDROCK_MODEL}/converse" + ), "the signature must cover the upstream URL the edge calls, not the edge URL the proxy called" + + def test_a_mount_without_a_signer_still_keys_on_its_credentials( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """The exclusion is per mount. Dropping authorization globally would let + one OpenAI account read another's recording.""" + cache: Final = bedrock_cache_edge(store) + assert "authorization" in SIGNATURE_HEADERS + assert "authorization" in cache.keyed("openai", HEADERS) + assert "authorization" not in cache.keyed(BEDROCK_MOUNT, HEADERS) + with edge(cache, provider) as url: + call(url) + with edge(bedrock_cache_edge(store), provider) as url: + call(url, headers=HEADERS | {"authorization": "Bearer synthetic-account-two"}) + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("action,response", [("converse", CONVERSE_SUCCESS), ("invoke", INVOKE_SUCCESS)]) + def test_complete_responses_replay_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("action,response", [ + ("converse", b'{"output":{"message":{}}}'), + ("converse", b'{"stopReason":"end_turn"}'), + ("converse", b'{"message":"The provided model identifier is invalid."}'), + ("converse", CONVERSE_SUCCESS[:-20]), + ("invoke", b'{"id":"msg_x","type":"message","content":[{"type":"text","text":"hi"}]}'), + ("invoke", b'{"id":"msg_x","type":"message","stop_reason":"end_turn"}'), + ("invoke", b'{"message":"Too many requests, please wait before trying again."}'), + ]) + def test_incomplete_or_error_bodies_never_enter_the_cache( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("action", ["converse-stream", "invoke-with-response-stream"]) + def test_streaming_endpoints_go_live_every_time( + self, store: RedisResponseStore, provider: Provider, action: str, + ) -> None: + """An eventstream's completeness cannot be proven without parsing its + frames, so these bypass rather than risk recording a truncated answer. + They are still signed: a bypass is a forward, not a passthrough.""" + provider.response = CONVERSE_SUCCESS + cache: Final = bedrock_cache_edge(store) + for _ in range(2): + with bedrock_edge(cache, provider, action) as url: + assert call(url, BEDROCK_BODY).body == CONVERSE_SUCCESS + assert len(provider.hits) == 2 + assert dict(cache.counters.counts)[f"mount:{BEDROCK_MOUNT}:bypass"] == 2 + assert all( + sent.startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") + for sent in provider.authorizations + ), provider.authorizations + + @pytest.mark.parametrize("action,cacheable", [ + ("converse", True), ("invoke", True), + ("converse-stream", False), ("invoke-with-response-stream", False), + ]) + def test_only_the_unary_bedrock_actions_are_cacheable(self, action: str, cacheable: bool) -> None: + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) is cacheable + + def test_a_region_mount_resolves_whole(self) -> None: + resolved: Final = resolve_mount(f"/{BEDROCK_MOUNT}/model/{BEDROCK_MODEL}/converse", EDGE_MOUNTS) + assert resolved is not None + assert resolved.mount == BEDROCK_MOUNT + assert resolved.upstream_base == "https://bedrock-runtime.us-east-1.amazonaws.com" + assert resolved.upstream_path == f"model/{BEDROCK_MODEL}/converse" + + def test_anthropic_stream_requires_start_finish_and_stop() -> None: start: Final = b'data: {"type":"message_start","message":{}}\n\n' finish: Final = b'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"}}\n\n' stop: Final = b'data: {"type":"message_stop"}\n\n' url: Final = "https://example.invalid/v1/messages" headers: Final = {"content-type": "text/event-stream"} - assert successful_response(url, 200, headers, start + finish + stop) - assert not successful_response(url, 200, headers, start + stop) - assert not successful_response(url, 200, headers, finish + stop) - assert not successful_response(url, 200, headers, start + finish) + assert successful_response("anthropic", url, 200, headers, start + finish + stop) + assert not successful_response("anthropic", url, 200, headers, start + stop) + assert not successful_response("anthropic", url, 200, headers, finish + stop) + assert not successful_response("anthropic", url, 200, headers, start + finish) @pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")]) @@ -342,6 +668,53 @@ def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMPa assert route_cache_model(params, unexpected_edge, enabled=True) is params +@pytest.mark.parametrize("model", [ + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "bedrock/converse/us.anthropic.claude-sonnet-5", + "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", +]) +def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str) -> None: + params: Final = LiteLLMParamsBody(model=model) + routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) + assert routed.aws_bedrock_runtime_endpoint == "http://edge.invalid/bedrock/us-east-1" + assert routed.api_base is None + assert routed.model_dump(exclude={"aws_bedrock_runtime_endpoint"}) == params.model_dump( + exclude={"aws_bedrock_runtime_endpoint"} + ) + + +@pytest.mark.parametrize("params", [ + LiteLLMParamsBody(model="bedrock/amazon.titan-embed-text-v2:0"), + LiteLLMParamsBody(model="bedrock/amazon.nova-canvas-v1:0"), + LiteLLMParamsBody(model="bedrock/amazon.nova-sonic-v1:0"), + LiteLLMParamsBody(model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_role_name="arn:aws:iam::1:role/x"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_access_key_id="AKIA"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", api_base="https://custom.invalid"), + LiteLLMParamsBody( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + aws_bedrock_runtime_endpoint="https://custom.invalid", + ), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_region_name="eu-west-1"), +]) +def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(params: LiteLLMParamsBody) -> None: + """Non-Anthropic models the runner role cannot invoke, deployments carrying + their own AWS identity (routing those would replace the assume-role chain the + batch suite exists to prove), explicit endpoints, and unmounted regions.""" + routed: Final = route_cache_model( + params, lambda mount: None if mount not in EDGE_MOUNTS else f"http://edge.invalid/{mount}", enabled=True, + ) + assert routed is params or routed.aws_bedrock_runtime_endpoint == params.aws_bedrock_runtime_endpoint + + +@pytest.mark.parametrize("mode", ["batch", "realtime", "image_generation"]) +def test_a_bedrock_deployment_with_a_mode_keeps_its_direct_route(mode: ModelMode) -> None: + params: Final = LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") + assert route_cache_model( + params, lambda mount: f"http://edge.invalid/{mount}", enabled=True, mode=mode, + ) is params + + def test_rollback_and_live_only_policy_keep_direct_provider_route() -> None: params: Final = LiteLLMParamsBody(model="openai/test") assert route_cache_model(params, lambda _: "http://edge.invalid", enabled=False) is params @@ -366,14 +739,16 @@ class PublishOutage: def test_write_outage_preserves_success_without_hidden_retry(store: RedisResponseStore, provider: Provider) -> None: unavailable: Final = replace(store, client=PublishOutage(store.client)) - cache: Final = CacheEdge(unavailable, SECRET) + cache: Final = cache_edge(unavailable) with edge(cache, provider) as url: assert call(url).body == SUCCESS assert call(url).body == SUCCESS assert len(provider.hits) == 2 assert dict(cache.counters.counts)["write_failures"] == 2 - with edge(CacheEdge(store, SECRET), provider) as url: + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS + assert len(provider.hits) == 3 + with edge(cache_edge(store), provider) as url: assert call(url).body == SUCCESS assert len(provider.hits) == 3 @@ -382,45 +757,41 @@ def test_connection_failure_releases_capture_lease(store: RedisResponseStore) -> with socket.socket() as unavailable: unavailable.bind(("127.0.0.1", 0)) url: Final = f"http://127.0.0.1:{unavailable.getsockname()[1]}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - assert isinstance(cache.forward("POST", url, HEADERS, BODY, 0.2), NetworkError) - prepared: Final = prepare_forward("POST", url, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY) - slot: Final = store.lookup(key) - assert isinstance(slot, CaptureLease) - assert store.release(key, slot) + cache: Final = cache_edge(store) + assert isinstance(cache.forward("openai", "POST", url, dict(HEADERS), BODY, 0.2), NetworkError) + key: Final = slot_key(url) + lease: Final = store.lookup(key) + assert isinstance(lease, CaptureLease) + assert store.release(key, lease) assert dict(cache.counters.counts)["rejected"] == 1 def test_close_before_first_chunk_releases_lease(store: RedisResponseStore, provider: Provider) -> None: url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - head: Final = cache.forward("POST", url, HEADERS, BODY, 5) + cache: Final = cache_edge(store) + head: Final = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) head.steps.close() - prepared: Final = prepare_forward("POST", url, HEADERS, BODY) - assert isinstance(prepared, PreparedForward) - key: Final = exact_key(SECRET, "POST", url, prepared.headers, BODY) - slot: Final = store.lookup(key) - assert isinstance(slot, CaptureLease) - assert store.release(key, slot) + key: Final = slot_key(url) + lease: Final = store.lookup(key) + assert isinstance(lease, CaptureLease) + assert store.release(key, lease) def test_effective_account_change_cannot_reuse_cache( store: RedisResponseStore, provider: Provider, monkeypatch: pytest.MonkeyPatch, tmp_path, ) -> None: url: Final = f"http://127.0.0.1:{provider.server_port}/v1/chat/completions" - cache: Final = CacheEdge(store, SECRET) - for account in ("account-a", "account-b", "account-b"): + caches: Final = tuple(cache_edge(store) for _ in range(3)) + for account, cache in zip(("account-a", "account-b", "account-b"), caches, strict=True): netrc = tmp_path / account netrc.write_text(f"machine 127.0.0.1 login {account} password synthetic\n") monkeypatch.setenv("NETRC", str(netrc)) - head = cache.forward("POST", url, HEADERS, BODY, 5) + head = cache.forward("openai", "POST", url, dict(HEADERS), BODY, 5) assert isinstance(head, StreamHead) assert b"".join(step.data for step in head.steps if isinstance(step, StreamChunk)) == SUCCESS assert len(provider.hits) == 2 - assert dict(cache.counters.counts)["hits"] == 1 + assert dict(caches[2].counters.counts)["hits"] == 1 def test_enabled_environment_reuses_store_across_fresh_backends( @@ -449,7 +820,7 @@ def test_enabled_environment_reuses_store_across_fresh_backends( def test_duplicate_headers_bypass_cache_and_count_live_calls( store: RedisResponseStore, provider: Provider, known_mount: bool, ) -> None: - cache: Final = CacheEdge(store, SECRET) + cache: Final = cache_edge(store) with edge(cache, provider) as url: parsed: Final = urlsplit(url) for _ in range(2): diff --git a/tests/e2e/fixture_canonical.py b/tests/e2e/fixture_canonical.py index e76d63ca33b..019c011aa67 100644 --- a/tests/e2e/fixture_canonical.py +++ b/tests/e2e/fixture_canonical.py @@ -51,6 +51,9 @@ SECRET_FIELD_SUFFIXES: Final[tuple[str, ...]] = ( ) SECRET_PLACEHOLDER: Final = "" +MARKER_PATTERN: Final = re.compile(r"(?" + PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = ( (re.compile(r"(?"), ( @@ -67,7 +70,7 @@ PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = ( re.compile(r"\b(?:chatcmpl|msgbatch|msg|resp|batch|call|req|ftjob|gen|file)[-_][A-Za-z0-9]{8,}\b"), "", ), - (re.compile(r"(?"), + (MARKER_PATTERN, MARKER_PLACEHOLDER), ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 7101438c5f8..7550bfdc150 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -951,6 +951,7 @@ class LiteLLMParamsBody(BaseModel): aws_access_key_id: str | None = None aws_secret_access_key: str | None = None aws_region_name: str | None = None + aws_bedrock_runtime_endpoint: str | None = None vertex_project: str | None = None vertex_location: str | None = None vertex_credentials: str | None = None diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 0c6eac75a43..1dc2f99abe5 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -23,12 +23,18 @@ from e2e_http import ( prepare_forward, primed_steps, ) +from fixture_canonical import MARKER_PATTERN, MARKER_PLACEHOLDER +from fixture_mode import SESSION_TEST_KEY, current_test_key from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError LIFETIME_SECONDS: Final = 86_400 MAX_REQUEST_BYTES: Final = 256 * 1024 MAX_RESPONSE_BYTES: Final = 8 * 1024 * 1024 UNRECORDED_RESPONSE_HEADERS: Final = frozenset({"set-cookie"}) +SIGNATURE_HEADERS: Final = frozenset( + {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} +) +BEDROCK_MOUNT_PREFIX: Final = "bedrock" JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -56,6 +62,7 @@ class CacheUnavailable: type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable +type RequestSigner = Callable[[str, str, Mapping[str, str], bytes | None], dict[str, str]] class ResponseStore(Protocol): @@ -83,28 +90,51 @@ class SignedResponse(BaseModel): signature: str -def exact_key(secret: bytes, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> str: +def canonical_text(value: str) -> str: + return MARKER_PATTERN.sub(MARKER_PLACEHOLDER, value) + + +def canonical_body(body: bytes) -> bytes: + try: + return canonical_text(body.decode("utf-8")).encode("utf-8") + except UnicodeDecodeError: + return body + + +def request_identity( + secret: bytes, test_key: str, method: str, url: str, headers: Mapping[str, str], body: bytes | None, +) -> str: fields: Final = ( - b"provider-cache-exact-v1", method.encode(), url.encode(), + b"provider-cache-canonical-v2", test_key.encode(), method.encode(), canonical_text(url).encode(), *(part.encode() for pair in sorted(headers.items()) for part in pair), - b"no-body" if body is None else b"body", b"" if body is None else body, + b"no-body" if body is None else b"body", b"" if body is None else canonical_body(body), ) encoded: Final = b"".join(len(part).to_bytes(8, "big") + part for part in fields) return hmac.new(secret, encoded, hashlib.sha256).hexdigest() -def cacheable_endpoint(method: str, url: str, body: bytes | None) -> bool: - return ( - method == "POST" - and urlsplit(url).path in {"/v1/chat/completions", "/v1/messages"} - and body is not None - and len(body) <= MAX_REQUEST_BYTES - ) +def slotted_key(secret: bytes, identity: str, slot: int) -> str: + return hmac.new(secret, f"{identity}:{slot}".encode(), hashlib.sha256).hexdigest() -def successful_response(url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool: +def is_bedrock(mount: str) -> bool: + return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX + + +def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool: + if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES: + return False + path: Final = urlsplit(url).path + if is_bedrock(mount): + return path.startswith("/model/") and path.endswith(("/converse", "/invoke")) + return path in {"/v1/chat/completions", "/v1/messages"} + + +def successful_response(mount: str, url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool: if not 200 <= status < 300 or len(body) > MAX_RESPONSE_BYTES: return False + if is_bedrock(mount): + return complete_bedrock_response(url, body) streaming: Final = "text/event-stream" in headers.get("content-type", "").lower() if streaming: try: @@ -147,6 +177,26 @@ def successful_response(url: str, status: int, headers: Mapping[str, str], body: ) +def complete_bedrock_response(url: str, body: bytes) -> bool: + """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an + Anthropic model answers the Anthropic message shape. Either way a truncated + or error body is missing the terminator field, which is what makes it safe to + record. The streaming variants never reach here: they are not cacheable.""" + try: + value: Final = JSON_VALUE.validate_json(body) + except ValidationError: + return False + if not isinstance(value, dict) or "message" in value: + return False + if urlsplit(url).path.endswith("/converse"): + return isinstance(value.get("output"), dict) and isinstance(value.get("stopReason"), str) + return ( + value.get("type") == "message" + and isinstance(value.get("content"), list) + and isinstance(value.get("stop_reason"), str) + ) + + def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool: if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values): return False @@ -172,7 +222,7 @@ def encode_response(secret: bytes, response: CachedResponse) -> bytes: return SignedResponse(response=raw, signature=hmac.new(secret, raw.encode(), hashlib.sha256).hexdigest()).model_dump_json().encode() -def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> CachedResponse | None: +def decode_response(secret: bytes, key: str, payload: bytes, mount: str, url: str) -> CachedResponse | None: if len(payload) > 2 * MAX_RESPONSE_BYTES: return None try: @@ -183,7 +233,9 @@ def decode_response(secret: bytes, key: str, payload: bytes, url: str) -> Cached chunks: Final = tuple(base64.b64decode(chunk, validate=True) for chunk in response.chunks) except (ValidationError, ValueError): return None - if response.request_key != key or not successful_response(url, response.status_code, response.headers, b"".join(chunks)): + if response.request_key != key or not successful_response( + mount, url, response.status_code, response.headers, b"".join(chunks) + ): return None return response @@ -199,6 +251,24 @@ class CacheCounters: self.counts = tuple((current | {name: current.get(name, 0) + 1}).items()) +@dataclass(slots=True) +class SlotCounter: + """FIFO position of a request among the canonically identical ones its test + has already sent. Two calls in one test that differ only by ``unique_marker`` + canonicalize the same, so without this they would share one recording and the + second would replay the first's provider response id.""" + + counts: tuple[tuple[str, int], ...] = () + lock: threading.Lock = field(default_factory=threading.Lock) + + def take(self, identity: str) -> int: + with self.lock: + current: Final = dict(self.counts) + taken: Final = current.get(identity, 0) + self.counts = tuple((current | {identity: taken + 1}).items()) + return taken + + @dataclass(slots=True) class ResponseCapture: buffer: io.BytesIO = field(default_factory=io.BytesIO) @@ -231,9 +301,12 @@ class CacheEdge: store: ResponseStore secret: bytes = field(repr=False) counters: CacheCounters = field(default_factory=CacheCounters) + slots: SlotCounter = field(default_factory=SlotCounter) + signers: Mapping[str, RequestSigner] = field(default_factory=dict) wait_seconds: float = 2.0 clock: Callable[[], float] = time.monotonic sleep: Callable[[float], None] = time.sleep + test_key: Callable[[], str] = current_test_key def lookup(self, key: str) -> CacheLookup: deadline: Final = self.clock() + self.wait_seconds @@ -241,39 +314,71 @@ class CacheEdge: self.sleep(min(0.05, max(0, deadline - self.clock()))) return result - def forward(self, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float) -> StreamHead | NetworkError: - if not cacheable_endpoint(method, url, body): - self.counters.increment("bypass") - self.counters.increment("upstream_attempts") - return forward_stream(method, url, headers=headers, body=body, timeout=timeout) - prepared: Final = prepare_forward(method, url, headers, body) + def count(self, mount: str, name: str) -> None: + self.counters.increment(name) + self.counters.increment(f"mount:{mount}:{name}") + + def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]: + """The headers actually sent upstream. A signing mount gets a signature + minted over the upstream URL, because the edge rewrote the Host the proxy + signed and Bedrock verifies it.""" + signer: Final = self.signers.get(mount) + return headers if signer is None else signer(method, url, headers, body) + + def keyed(self, mount: str, headers: Mapping[str, str]) -> Mapping[str, str]: + """A signing mount's signature headers are the edge's own and carry a + timestamp, so keying on them would make every request a permanent miss. + Every other mount keys on its headers whole, credentials included, so a + different account can never read another's recording.""" + if mount not in self.signers: + return headers + return {name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS} + + def forward( + self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float, + ) -> StreamHead | NetworkError: + test_key: Final = self.test_key() + if test_key == SESSION_TEST_KEY or not cacheable_endpoint(mount, method, url, body): + self.count(mount, "bypass") + self.count(mount, "upstream_attempts") + return forward_stream( + method, url, headers=self.outbound(mount, method, url, headers, body), body=body, timeout=timeout, + ) + prepared: Final = prepare_forward(method, url, self.outbound(mount, method, url, headers, body), body) if isinstance(prepared, NetworkError): - self.counters.increment("rejected") + self.count(mount, "rejected") return prepared - key: Final = exact_key(self.secret, method, url, prepared.headers, body) + identity: Final = request_identity( + self.secret, test_key, method, url, self.keyed(mount, prepared.headers), body, + ) + key: Final = slotted_key(self.secret, identity, self.slots.take(identity)) found: Final = self.lookup(key) if isinstance(found, CacheHit): - response: Final = decode_response(self.secret, key, found.payload, url) + response: Final = decode_response(self.secret, key, found.payload, mount, url) if response is not None and self.clock() < found.valid_until: - self.counters.increment("hits") + self.count(mount, "hits") return StreamHead(response.status_code, response.headers, response_steps(response)) - self.counters.increment("corrupt" if response is None else "expired") + self.count(mount, "corrupt" if response is None else "expired") self.store.discard(key, found.payload) capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found - self.counters.increment("misses") + self.count(mount, "misses") if isinstance(capture_slot, CacheUnavailable): - self.counters.increment("cache_errors") - self.counters.increment("upstream_attempts") + self.count(mount, "cache_errors") + self.count(mount, "upstream_attempts") head: Final = forward_prepared_stream(prepared, timeout) if not isinstance(capture_slot, CaptureLease): return head if isinstance(head, NetworkError): self.store.release(key, capture_slot) - self.counters.increment("rejected") + self.count(mount, "rejected") return head - return StreamHead(head.status_code, head.headers, primed_steps(self.capture(key, capture_slot, url, head))) + return StreamHead( + head.status_code, head.headers, primed_steps(self.capture(mount, key, capture_slot, url, head)), + ) - def capture(self, key: str, lease: CaptureLease, url: str, head: StreamHead) -> Generator[StreamStep, None, None]: + def capture( + self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, + ) -> Generator[StreamStep, None, None]: capture: Final = ResponseCapture() try: with closing(head.steps): @@ -285,15 +390,15 @@ class CacheEdge: headers: Final = { name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS } - if not capture.eligible or not successful_response(url, head.status_code, headers, b"".join(chunks)): - self.counters.increment("rejected") + if not capture.eligible or not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): + self.count(mount, "rejected") return response: Final = CachedResponse( request_key=key, status_code=head.status_code, headers=headers, chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks), ) published: Final = self.store.publish(key, lease, encode_response(self.secret, response)) - self.counters.increment("writes" if published else "write_failures") + self.count(mount, "writes" if published else "write_failures") finally: self.store.release(key, lease) capture.buffer.close() diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index 24599b5a313..e7e4899eb71 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -8,14 +8,57 @@ from models import LiteLLMParamsBody, ModelMode LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False) +DEFAULT_BEDROCK_REGION: Final = "us-east-1" +BEDROCK_ANTHROPIC_INFIX: Final = "anthropic." + + +def bedrock_mount(params: LiteLLMParamsBody) -> str | None: + """The edge mount an Anthropic-on-Bedrock deployment belongs to, or None. + + Only the Anthropic models route. The edge validates converse and invoke + bodies by their Anthropic and Converse terminator fields, and the runner role + is allowed to invoke exactly those models, so Bedrock embeddings, image + generation, rerank and realtime keep their existing direct path rather than + reaching an edge that could neither sign nor validate for them.""" + route: Final = params.model.partition("/")[2] + model: Final = route.partition("/")[2] or route + if BEDROCK_ANTHROPIC_INFIX not in model: + return None + return f"bedrock/{params.aws_region_name or DEFAULT_BEDROCK_REGION}" + + +def route_bedrock( + params: LiteLLMParamsBody, base_for: Callable[[str], str | None], mode: ModelMode | None, +) -> LiteLLMParamsBody: + """Deployments that carry their own AWS identity stay off the edge. The edge + re-signs with the run pod's role, so routing an `aws_role_name` deployment + would quietly replace the very assume-role chain that test exists to prove.""" + if mode is not None or params.aws_role_name is not None or params.aws_access_key_id is not None: + return params + if params.api_base is not None or params.aws_bedrock_runtime_endpoint is not None: + return params + mount: Final = bedrock_mount(params) + if mount is None: + return params + base: Final = base_for(mount) + if base is None: + return params + return params.model_copy(update={"aws_bedrock_runtime_endpoint": base}) + def route_cache_model( params: LiteLLMParamsBody, base_for: Callable[[str], str | None], *, enabled: bool, mode: ModelMode | None = None, ) -> LiteLLMParamsBody: - if not enabled or mode == "realtime" or LIVE_PROVIDER_REQUIRED.get() or params.api_base is not None or params.mock_response is not None: + if not enabled or LIVE_PROVIDER_REQUIRED.get() or params.mock_response is not None: + return params + if params.litellm_credential_name is not None: return params provider: Final = params.model.partition("/")[0] - if provider not in {"openai", "anthropic"} or params.litellm_credential_name is not None: + if provider == "bedrock": + return route_bedrock(params, base_for, mode) + if mode == "realtime" or params.api_base is not None: + return params + if provider not in {"openai", "anthropic"}: return params base: Final = base_for(provider) if base is None: diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index dda9e6f8e4f..8f718ad1967 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -48,7 +48,7 @@ import threading from collections import deque from collections.abc import Generator, Mapping, Sequence from contextlib import closing, contextmanager -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import islice from pathlib import Path @@ -94,17 +94,41 @@ from fixture_mode import ( parse_fixture_mode, ) from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity -from provider_cache import CacheEdge +from provider_cache import CacheEdge, RequestSigner, is_bedrock from provider_cache_routing import LIVE_PROVIDER_REQUIRED from pydantic import JsonValue, TypeAdapter +BEDROCK_REGIONS: Final[tuple[str, ...]] = ("us-east-1",) + EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { "openai": "https://api.openai.com", "anthropic": "https://api.anthropic.com", + **{ + f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com" + for region in BEDROCK_REGIONS + }, } ) + +@dataclass(frozen=True, slots=True) +class ResolvedMount: + mount: str + upstream_base: str + upstream_path: str + + +def resolve_mount(path: str, mounts: Mapping[str, str]) -> ResolvedMount | None: + """Longest mount prefix wins, so a region-qualified mount such as + ``bedrock/us-east-1`` resolves whole instead of leaving the region as the + first segment of the upstream path.""" + trimmed: Final = path.lstrip("/") + for mount in sorted(mounts, key=len, reverse=True): + if trimmed == mount or trimmed.startswith(f"{mount}/"): + return ResolvedMount(mount, mounts[mount], trimmed[len(mount):].lstrip("/")) + return None + REPLAY_MISS_STATUS: Final = 599 _HOP_BY_HOP_HEADERS: Final[frozenset[str]] = frozenset( @@ -754,14 +778,14 @@ def _handle_record( def _handle_live( method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float, - cache: CacheEdge | None = None, + cache: CacheEdge | None = None, mount: str = "", ) -> EdgeOutcome: forwarded: Final = { name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS } head: Final = ( forward_stream(method, url, headers=forwarded, body=body, timeout=timeout) - if cache is None else cache.forward(method, url, forwarded, body, timeout) + if cache is None else cache.forward(mount, method, url, forwarded, body, timeout) ) match head: case NetworkError(message=message): @@ -796,10 +820,13 @@ def handle_edge_request( prefix, then record (forward + persist) or replay (serve from the bundle). Socket-free so unit tests exercise every branch without a server.""" split: Final = urlsplit(raw_path) - mount, _, upstream_path = split.path.lstrip("/").partition("/") - upstream_base: Final = mounts.get(mount) - if upstream_base is None: - return _text_reply(404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}") + resolved: Final = resolve_mount(split.path, mounts) + if resolved is None: + unknown: Final = split.path.lstrip("/").partition("/")[0] + return _text_reply(404, f"unknown provider mount {unknown!r}; known mounts: {', '.join(sorted(mounts))}") + mount: Final = resolved.mount + upstream_base: Final = resolved.upstream_base + upstream_path: Final = resolved.upstream_path profile: Final = ( backend.recorder.profile if isinstance(backend, RecordEdge) @@ -830,7 +857,8 @@ def handle_edge_request( match backend: case CacheEdge(): return _handle_live( - method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, backend, + method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout, + backend, mount, ) case LiveEdge(): return _handle_live( @@ -891,7 +919,7 @@ class _EdgeHandler(BaseHTTPRequestHandler): ) if isinstance(edge_server.backend, CacheEdge) and duplicate_headers: edge_server.backend.counters.increment("duplicate_header_bypass") - if urlsplit(self.path).path.lstrip("/").partition("/")[0] in edge_server.mounts: + if resolve_mount(urlsplit(self.path).path, edge_server.mounts) is not None: edge_server.backend.counters.increment("upstream_attempts") outcome: Final = handle_edge_request( selected_backend, @@ -1079,6 +1107,8 @@ def provider_edge_api_base( return _shared_cache_edge(bind_host, advertise_host, forward_timeout).api_base(mount) return None case "record" | "replay": + if is_bedrock(mount): + return None if mount not in EDGE_MOUNTS: raise ValueError(f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(EDGE_MOUNTS))}") return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout, match_profile()).api_base( @@ -1108,7 +1138,17 @@ def configured_cache_backend() -> CacheEdge | None: return None from provider_cache_redis import configured_cache - return configured_cache() + cache: Final = configured_cache() + return None if cache is None else replace(cache, signers=bedrock_signers()) + + +@functools.lru_cache(maxsize=1) +def bedrock_signers() -> Mapping[str, RequestSigner]: + """One signer per mounted Bedrock region, built lazily so a run that never + mounts Bedrock neither imports botocore nor resolves an AWS identity.""" + from provider_edge_bedrock import bedrock_signer + + return MappingProxyType({f"bedrock/{region}": bedrock_signer(region) for region in BEDROCK_REGIONS}) @functools.lru_cache(maxsize=8) diff --git a/tests/e2e/provider_edge_bedrock.py b/tests/e2e/provider_edge_bedrock.py new file mode 100644 index 00000000000..5d8148482d7 --- /dev/null +++ b/tests/e2e/provider_edge_bedrock.py @@ -0,0 +1,72 @@ +"""SigV4 re-signing for Bedrock traffic routed through the provider edge. + +Bedrock is the one provider the edge could never mount. SigV4 signs the Host +header, so rewriting ``api_base`` to point at the edge invalidates the proxy's +signature and Bedrock rejects the call before it reaches a model. The edge +therefore has to drop the proxy's signature and mint its own over the upstream +URL it is actually about to call. + +The identity it signs with is the run pod's own, from the EKS Pod Identity +association on ServiceAccount ``buildkite-e2e-run``. That role carries Bedrock +invoke and converse on an allowlist of the Anthropic models the suite registers +and nothing else, so a re-signed call can reach exactly the models the suite +already uses. The proxy's own Bedrock credentials are not involved in a routed +deployment, which is why ``aws_role_name`` deployments stay off the edge: their +whole point is to prove the product's assume-role chain. + +Signature headers are excluded from the cache key by the caller, and they have +to be: ``x-amz-date`` is a timestamp, so keying on it would make every Bedrock +request a permanent miss. +""" + +from __future__ import annotations + +import functools +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials +from botocore.session import Session +from provider_cache import SIGNATURE_HEADERS + +BEDROCK_SERVICE: Final = "bedrock" + + +class MissingAwsCredentials(RuntimeError): + """No AWS identity is resolvable, so the edge cannot sign for Bedrock.""" + + +@dataclass(frozen=True, slots=True) +class BedrockSigner: + region: str + credentials: Callable[[], Credentials] + + def __call__(self, method: str, url: str, headers: Mapping[str, str], body: bytes | None) -> dict[str, str]: + unsigned: Final = { + name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS + } + request: Final = AWSRequest(method=method, url=url, headers=unsigned, data=body or b"") + SigV4Auth(self.credentials(), BEDROCK_SERVICE, self.region).add_auth(request) + return dict(request.headers) + + +@functools.lru_cache(maxsize=1) +def pod_credentials() -> Credentials: + """The run pod's own identity, resolved once per process through botocore's + ordinary chain, which reaches Pod Identity at the ``container-role`` link.""" + resolved: Final = Session().get_credentials() + if resolved is None: + raise MissingAwsCredentials( + "the provider edge is mounted for Bedrock but no AWS credentials resolve; " + "the run pod gets them from the Pod Identity association on buildkite-e2e-run" + ) + return resolved + + +def bedrock_signer(region: str, credentials: Callable[[], Credentials] = pod_credentials) -> BedrockSigner: + """Credentials are resolved on the first signed request, not here, so a run + that mounts Bedrock but never calls it needs no AWS identity at all.""" + return BedrockSigner(region, credentials) diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index 5d0c79f26f6..c8d70697182 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -1279,15 +1279,30 @@ class TestApiBaseSeam: ) def test_unknown_mount_raises_naming_the_known_mounts(self, tmp_path: Path) -> None: - with pytest.raises(ValueError, match="unknown provider mount 'bedrock'"): + with pytest.raises(ValueError, match="unknown provider mount 'cohere'"): provider_edge_api_base( - "bedrock", + "cohere", mode_raw="record", bundle_dir=tmp_path / "bundle", bind_host="127.0.0.1", advertise_host="127.0.0.1", ) + @pytest.mark.parametrize("mode_raw", ["record", "replay"]) + def test_bedrock_never_wires_a_bundle_because_the_edge_cannot_sign_into_one( + self, tmp_path: Path, mode_raw: str, + ) -> None: + """Record and replay serve from a bundle without re-signing, so a Bedrock + deployment pointed at that edge would send the proxy's signature over a + rewritten Host. It keeps its direct route in both modes.""" + assert provider_edge_api_base( + "bedrock/us-east-1", + mode_raw=mode_raw, + bundle_dir=tmp_path / "bundle", + bind_host="127.0.0.1", + advertise_host="127.0.0.1", + ) is None + def test_record_mode_boots_one_shared_edge_and_prepares_the_bundle(self, tmp_path: Path) -> None: root = tmp_path / "bundle" first = provider_edge_api_base( From b68e60f7061a2102fa076ff7ca6368f819d13770 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 02:34:09 -0700 Subject: [PATCH 071/187] feat(e2e): cache the responses and embeddings endpoints behind the edge Chat completions and messages were the only cacheable paths. The suite also drives /v1/embeddings and /v1/responses through the same OpenAI mount, so both now cache, each with its own completeness rule: a chat response's `choices` check would reject a perfectly good embedding, and a Responses run that never reached `response.completed` must stay out of the cache the same way a truncated stream does. Vertex and Gemini stay off the edge. litellm's `_check_custom_proxy` rewrites a path-prefixed vertex api_base into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without a root-mounted edge on its own port or a change in litellm. Shipping an unvalidated URL guess would have been worse than saying so in PROVIDER_CACHE.md. Also finishes the MountPolicy move: a mount now carries its signer and its unkeyed headers together instead of a bare signer map. --- .../test_provider_cache.py | 102 +++++++++++++++++- tests/e2e/PROVIDER_CACHE.md | 28 ++++- tests/e2e/provider_cache.py | 63 +++++++++-- tests/e2e/provider_edge.py | 15 ++- tests/e2e/provider_edge_bedrock.py | 2 +- 5 files changed, 186 insertions(+), 24 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 5d35981344d..5c491e1cdcd 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -25,6 +25,7 @@ from provider_cache import ( CacheEdge, CacheHit, CaptureLease, + MountPolicy, ResponseStore, cacheable_endpoint, request_identity, @@ -179,7 +180,9 @@ def slot_key( def bedrock_cache_edge(store: ResponseStore, test_key: str = TEST_KEY) -> CacheEdge: return CacheEdge( store, SECRET, test_key=lambda: test_key, - signers={BEDROCK_MOUNT: bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS)}, + policies={BEDROCK_MOUNT: MountPolicy( + sign=bedrock_signer("us-east-1", lambda: STATIC_CREDENTIALS), unkeyed_headers=SIGNATURE_HEADERS, + )}, ) @@ -501,6 +504,98 @@ def test_counters_attribute_every_outcome_to_its_mount( assert counts["mount:anthropic:rejected"] == 1 and "mount:openai:rejected" not in counts +EMBEDDING_SUCCESS: Final = ( + b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}],' + b'"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' +) +RESPONSE_SUCCESS: Final = b'{"id":"resp_synthetic","object":"response","status":"completed","output":[]}' +RESPONSE_STREAM_SUCCESS: Final = ( + b'data: {"type":"response.created","response":{"id":"resp_synthetic"}}\n\n' + b'data: {"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"}}\n\n' +) + + +@contextmanager +def openai_edge(cache: CacheEdge, provider: Provider, path: str) -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + running: Final = start_provider_edge(cache, mounts={"openai": upstream}) + try: + yield running.edge.api_base("openai") + path + finally: + running.shutdown() + + +class TestNonChatOpenAiEndpoints: + """Chat and messages were the only cacheable paths. Embeddings and responses + are the other two JSON endpoints the suite drives through the same mount, and + each needs its own completeness rule: a chat response's ``choices`` check + would reject a perfectly good embedding.""" + + @pytest.mark.parametrize("path,response", [ + ("/v1/embeddings", EMBEDDING_SUCCESS), + ("/v1/responses", RESPONSE_SUCCESS), + ]) + def test_complete_responses_replay_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + def test_a_completed_response_stream_replays( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + provider.stream = True + provider.response = RESPONSE_STREAM_SUCCESS + for _ in range(2): + with openai_edge(cache_edge(store), provider, "/v1/responses") as url: + assert call(url, MARKED).body == RESPONSE_STREAM_SUCCESS + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("path,response", [ + ("/v1/embeddings", b'{"object":"list","data":[],"usage":{"prompt_tokens":0}}'), + ("/v1/embeddings", b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[]}],"usage":{}}'), + ("/v1/embeddings", b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1]}]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","status":"incomplete","output":[]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","status":"in_progress","output":[]}'), + ("/v1/responses", b'{"id":"resp_x","object":"response","output":[]}'), + ]) + def test_incomplete_bodies_never_enter_the_cache( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("payload", [ + b'data: {"type":"response.created","response":{"id":"resp_x"}}\n\n', + b'data: {"type":"response.created","response":{"id":"resp_x"}}\n\ndata: {"type":"response.failed"}\n\n', + b'data: {"type":"response.completed","response":{"id":"resp_x"}}\n\ndata: {"type":"response.created"}\n\n', + ]) + def test_a_response_stream_that_never_completed_is_never_cached( + self, store: RedisResponseStore, provider: Provider, payload: bytes, + ) -> None: + provider.stream = True + provider.response = payload + for _ in range(2): + with openai_edge(cache_edge(store), provider, "/v1/responses") as url: + assert call(url, MARKED).body == payload + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("path,cacheable", [ + ("/v1/chat/completions", True), ("/v1/messages", True), + ("/v1/embeddings", True), ("/v1/responses", True), + ("/v1/audio/speech", False), ("/v1/images/generations", False), + ("/v1/files", False), ("/v1/batches", False), + ]) + def test_only_the_json_endpoints_are_cacheable(self, path: str, cacheable: bool) -> None: + assert cacheable_endpoint("openai", "POST", f"https://api.openai.com{path}", MARKED) is cacheable + + class TestBedrockSigning: """Bedrock is the reason the edge could not mount it before: SigV4 covers the Host header, so forwarding through a rewritten api_base invalidates the @@ -545,7 +640,10 @@ class TestBedrockSigning: return dict(headers) | {"authorization": f"AWS4-HMAC-SHA256 {url}", "x-amz-date": next(stamps)} def signing_edge() -> CacheEdge: - return CacheEdge(store, SECRET, test_key=lambda: TEST_KEY, signers={BEDROCK_MOUNT: varying}) + return CacheEdge( + store, SECRET, test_key=lambda: TEST_KEY, + policies={BEDROCK_MOUNT: MountPolicy(sign=varying, unkeyed_headers=SIGNATURE_HEADERS)}, + ) for _ in range(2): with bedrock_edge(signing_edge(), provider) as url: diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 8635c9ed9ae..393a96e8c16 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -1,11 +1,29 @@ # Shared provider-response cache -`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live +`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live -The edge caches complete successful POST responses for `/v1/chat/completions` and `/v1/messages`, including streams. Unsupported endpoints pass through. It matches the method, original URL, effective outbound headers (including authentication and HTTP-library defaults), body presence and exact body bytes using a full keyed digest. It sends the same prepared request used for matching. No prompts, random markers, JSON values or credentials are normalized away. Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, including streams. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored + +## Request identity + +A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is + +Requests that differ only by their markers therefore share a canonical identity, which is what makes the cache reusable across builds: every e2e test salts its prompt afresh, so an exact-byte key would miss on every call. Within one test, calls that share a canonical identity are still recorded and replayed separately, by a FIFO slot index appended to the key. That matters because a replayed response carries the recorded provider response id, `LiteLLM_SpendLogs.request_id` is that id, and one shared recording answering two calls would collapse two spend rows into one + +Two different tests never share a recording, and a provider call made outside any test (fixtures, session setup) is never cached, because the identity has no test node id to bind to + +Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure +## Bedrock + +Bedrock could not be mounted before because SigV4 signs the `Host` header, so a rewritten `api_base` failed signature verification at the provider. The edge now re-signs: it drops the proxy's signature headers, signs the upstream request with the run pod's own AWS identity from its EKS Pod Identity association, and forwards that. The signature headers are excluded from the key, since `x-amz-date` is a timestamp and keying on it would make every Bedrock call a permanent miss + +Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove. Only Anthropic models route, matching what the runner role is allowed to invoke and what the edge knows how to validate + +Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm + Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires ## Configuration @@ -18,16 +36,16 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay ## Recorded response semantics -Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Existing spend reconciliation requests use distinct prompt markers and retain their distinct-ID and row-count assertions; accounting tests are not automatically excluded from caching +Replay preserves the original response ID, usage and end-to-end headers. The proxy can therefore deduplicate repeated provider IDs when storing spend-log rows, just as it does when a live upstream returns the same ID twice. One spend-log row per invocation is not guaranteed for identical recorded responses. Spend reconciliation keeps its distinct-ID and row-count assertions: its prompts differ by an index as well as a marker, so they stay distinct once markers are normalized, and calls that are canonically equal within one test take separate FIFO slots and separate recordings anyway. Accounting tests are not automatically excluded from caching Provider remaining-quota headers describe the captured response. Metrics derived from them are historical on a cache hit, not a measurement of current provider capacity. Gateway-generated API-key quota headers are a separate contract. A test of fresh provider quota or timing must use the live-provider policy; replay can still exercise how the proxy processes the recorded headers ## Qualification -`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence +`tests/code_coverage_tests/test_provider_cache.py` exercises local HTTP providers and disposable real Redis, including the marker-canonical key, the FIFO slot index, per-test isolation, SigV4 re-signing against a local upstream, and each endpoint's completeness rule. CI runs these checks with the existing provider-edge and replay harness tests. These component checks do not establish Buildkite deployment, full-suite cross-build reuse or a genuine 24-hour expiry observation; those require separate runtime evidence diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 1dc2f99abe5..a2b18a3a466 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -9,6 +9,7 @@ import time from collections.abc import Callable, Generator, Mapping from contextlib import closing from dataclasses import dataclass, field +from types import MappingProxyType from typing import Final, Literal, Protocol from urllib.parse import urlsplit @@ -35,6 +36,7 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" +OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -65,6 +67,23 @@ type CacheLookup = CacheHit | CaptureLease | CacheBusy | CacheUnavailable type RequestSigner = Callable[[str, str, Mapping[str, str], bytes | None], dict[str, str]] +@dataclass(frozen=True, slots=True) +class MountPolicy: + """What a mount needs beyond plain forwarding. + + ``sign`` mints a fresh credential over the upstream URL, for providers whose + auth covers the Host the edge rewrote. ``unkeyed_headers`` names headers that + must stay out of the cache key because they change on every call and would + otherwise make the mount a permanent miss: a minted signature, or an OAuth + token the provider rotates. Naming one costs the guarantee that a recording + can never cross credentials, so a mount with a rotating token relies on the + environment holding one identity for that provider. Mounts with a static API + key name nothing here and keep the guarantee whole.""" + + sign: RequestSigner | None = None + unkeyed_headers: frozenset[str] = frozenset() + + class ResponseStore(Protocol): def lookup(self, key: str) -> CacheLookup: ... @@ -127,7 +146,7 @@ def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> path: Final = urlsplit(url).path if is_bedrock(mount): return path.startswith("/model/") and path.endswith(("/converse", "/invoke")) - return path in {"/v1/chat/completions", "/v1/messages"} + return path in OPENAI_JSON_PATHS def successful_response(mount: str, url: str, status: int, headers: Mapping[str, str], body: bytes) -> bool: @@ -150,6 +169,8 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values): return False + if urlsplit(url).path == "/v1/responses": + return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values) return ( @@ -168,8 +189,17 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not isinstance(value, dict) or "error" in value: return False - if urlsplit(url).path == "/v1/messages": + path: Final = urlsplit(url).path + if path == "/v1/messages": return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str) + if path == "/v1/embeddings": + data: Final = value.get("data") + return isinstance(data, list) and bool(data) and isinstance(value.get("usage"), dict) and all( + isinstance(item, dict) and isinstance(item.get("embedding"), list) and bool(item["embedding"]) + for item in data + ) + if path == "/v1/responses": + return value.get("object") == "response" and value.get("status") == "completed" choices: Final = value.get("choices") return isinstance(choices, list) and bool(choices) and all( isinstance(choice, dict) and isinstance(choice.get("message"), dict) and isinstance(choice.get("finish_reason"), str) @@ -197,6 +227,14 @@ def complete_bedrock_response(url: str, body: bytes) -> bool: ) +def complete_responses_stream(values: tuple[JsonValue, ...]) -> bool: + """The Responses API streams typed events and ends with ``response.completed``. + A run that failed, was cancelled, or ran out of tokens ends with a different + terminal event, so requiring that one keeps a half-finished response out.""" + last: Final = values[-1] + return isinstance(last, dict) and last.get("type") == "response.completed" + + def complete_chat_stream(values: tuple[JsonValue, ...]) -> bool: if any(not isinstance(value, dict) or not isinstance(value.get("choices"), list) for value in values): return False @@ -296,13 +334,16 @@ def response_steps(response: CachedResponse) -> Generator[StreamStep, None, None yield StreamChunk(base64.b64decode(chunk, validate=True)) +NO_POLICIES: Final[Mapping[str, MountPolicy]] = MappingProxyType({}) + + @dataclass(frozen=True, slots=True) class CacheEdge: store: ResponseStore secret: bytes = field(repr=False) counters: CacheCounters = field(default_factory=CacheCounters) slots: SlotCounter = field(default_factory=SlotCounter) - signers: Mapping[str, RequestSigner] = field(default_factory=dict) + policies: Mapping[str, MountPolicy] = NO_POLICIES wait_seconds: float = 2.0 clock: Callable[[], float] = time.monotonic sleep: Callable[[float], None] = time.sleep @@ -321,18 +362,18 @@ class CacheEdge: def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]: """The headers actually sent upstream. A signing mount gets a signature minted over the upstream URL, because the edge rewrote the Host the proxy - signed and Bedrock verifies it.""" - signer: Final = self.signers.get(mount) + signed and the provider verifies it.""" + signer: Final = self.policies.get(mount, MountPolicy()).sign return headers if signer is None else signer(method, url, headers, body) def keyed(self, mount: str, headers: Mapping[str, str]) -> Mapping[str, str]: - """A signing mount's signature headers are the edge's own and carry a - timestamp, so keying on them would make every request a permanent miss. - Every other mount keys on its headers whole, credentials included, so a - different account can never read another's recording.""" - if mount not in self.signers: + """Headers the cache key is built from. A mount keeps its credentials in + the key unless its policy names them unkeyed, so by default one account + can never read another's recording.""" + unkeyed: Final = self.policies.get(mount, MountPolicy()).unkeyed_headers + if not unkeyed: return headers - return {name: value for name, value in headers.items() if name.lower() not in SIGNATURE_HEADERS} + return {name: value for name, value in headers.items() if name.lower() not in unkeyed} def forward( self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None, timeout: float, diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 8f718ad1967..2606b26fe99 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -94,7 +94,7 @@ from fixture_mode import ( parse_fixture_mode, ) from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity -from provider_cache import CacheEdge, RequestSigner, is_bedrock +from provider_cache import SIGNATURE_HEADERS, CacheEdge, MountPolicy, is_bedrock from provider_cache_routing import LIVE_PROVIDER_REQUIRED from pydantic import JsonValue, TypeAdapter @@ -1139,16 +1139,21 @@ def configured_cache_backend() -> CacheEdge | None: from provider_cache_redis import configured_cache cache: Final = configured_cache() - return None if cache is None else replace(cache, signers=bedrock_signers()) + return None if cache is None else replace(cache, policies=bedrock_policies()) @functools.lru_cache(maxsize=1) -def bedrock_signers() -> Mapping[str, RequestSigner]: - """One signer per mounted Bedrock region, built lazily so a run that never +def bedrock_policies() -> Mapping[str, MountPolicy]: + """One policy per mounted Bedrock region, built lazily so a run that never mounts Bedrock neither imports botocore nor resolves an AWS identity.""" from provider_edge_bedrock import bedrock_signer - return MappingProxyType({f"bedrock/{region}": bedrock_signer(region) for region in BEDROCK_REGIONS}) + return MappingProxyType( + { + f"bedrock/{region}": MountPolicy(sign=bedrock_signer(region), unkeyed_headers=SIGNATURE_HEADERS) + for region in BEDROCK_REGIONS + } + ) @functools.lru_cache(maxsize=8) diff --git a/tests/e2e/provider_edge_bedrock.py b/tests/e2e/provider_edge_bedrock.py index 5d8148482d7..73e4a16d272 100644 --- a/tests/e2e/provider_edge_bedrock.py +++ b/tests/e2e/provider_edge_bedrock.py @@ -58,7 +58,7 @@ def pod_credentials() -> Credentials: """The run pod's own identity, resolved once per process through botocore's ordinary chain, which reaches Pod Identity at the ``container-role`` link.""" resolved: Final = Session().get_credentials() - if resolved is None: + if resolved is None: # pyright: ignore[reportUnnecessaryComparison] # stubs miss the empty-chain None raise MissingAwsCredentials( "the provider edge is mounted for Bedrock but no AWS credentials resolve; " "the run pod gets them from the Pod Identity association on buildkite-e2e-run" From 6cf35ed71bf769154eb70bfff69e6c70202961e6 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 09:42:20 +0000 Subject: [PATCH 072/187] feat(proxy): expose lifetime total_spend on virtual keys Adds a persistent total_spend column to LiteLLM_VerificationToken and LiteLLM_DeletedVerificationToken, incremented in the same write as spend and left alone by budget resets. Surfaces it on /key/info, /key/list and the Admin UI Virtual Keys table and key detail view Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 5 ++ .../litellm_proxy_extras/schema.prisma | 2 + litellm/models/verification_token.py | 1 + litellm/proxy/db/db_spend_update_writer.py | 1 + litellm/proxy/schema.prisma | 2 + schema.prisma | 2 + tests/test_litellm/models/test_models.py | 8 +++ .../common_utils/test_reset_budget_job.py | 17 ++++++ .../proxy/db/test_db_spend_update_writer.py | 54 ++++++++++++++++++- .../test_key_management_endpoints.py | 54 +++++++++++++++++++ .../DeletedKeysPage/DeletedKeysPage.test.tsx | 1 + .../VirtualKeysPage/VirtualKeysTable.test.tsx | 9 ++++ .../VirtualKeysPage/keyTableColumns.tsx | 15 ++++++ .../components/key_team_helpers/key_list.tsx | 1 + .../key_edit_view.integration.test.tsx | 1 + .../templates/key_info_view.test.tsx | 18 +++++++ .../components/templates/key_info_view.tsx | 8 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 15 ++++++ 18 files changed, 213 insertions(+), 1 deletion(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql new file mode 100644 index 00000000000..daacd66db39 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; + +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d2375903c47..139fb031671 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -426,6 +426,7 @@ model LiteLLM_VerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") @@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index 06ff877a41a..15ecf5fe026 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -18,6 +18,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): key_name: str | None = None key_alias: str | None = None spend: float = 0.0 + total_spend: float = 0.0 max_budget: float | None = None expires: str | datetime | None = None models: list = [] diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index a90d1351fd7..5b43ed53117 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1619,6 +1619,7 @@ class DBSpendUpdateWriter: where={"token": token}, data={ "spend": {"increment": response_cost}, + "total_spend": {"increment": response_cost}, "last_active": datetime.now(timezone.utc), }, ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d2375903c47..139fb031671 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -426,6 +426,7 @@ model LiteLLM_VerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") @@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") diff --git a/schema.prisma b/schema.prisma index d2375903c47..139fb031671 100644 --- a/schema.prisma +++ b/schema.prisma @@ -426,6 +426,7 @@ model LiteLLM_VerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") @@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index aa6449c98dd..9b803c14062 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -362,6 +362,14 @@ class TestVerificationToken: assert deleted.deleted_at is not None assert deleted.token == "t1" + def test_total_spend_is_carried_separately_from_resettable_spend(self): + token = LiteLLM_VerificationToken(token="t1", spend=0.0, total_spend=12.5) + assert token.model_dump()["total_spend"] == 12.5 + assert token.model_dump()["spend"] == 0.0 + + deleted = LiteLLM_DeletedVerificationToken.model_validate({**token.model_dump(), "deleted_by": "admin"}) + assert deleted.total_spend == 12.5 + class TestConfigTable: def test_config_creation(self): diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 943a6c905c0..1ccf9be37b9 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -291,6 +291,23 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): assert set(write["data"].keys()) == {"spend", "budget_reset_at"} +def test_reset_budget_for_key_leaves_lifetime_total_spend_alone(reset_budget_job, mock_prisma_client): + """A period reset zeroes spend but must neither write nor touch the lifetime total_spend.""" + now = datetime.now(timezone.utc) + key = LiteLLM_VerificationToken( + token="tok-key-1", spend=100.0, total_spend=340.0, budget_duration="30d", budget_reset_at=now + ) + mock_prisma_client.data["key"] = [key] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + (write,) = _batch_writes(mock_prisma_client, "key") + assert write["data"]["spend"] == {"decrement": 100.0} + assert "total_spend" not in write["data"] + assert key.spend == 0.0 + assert key.total_spend == 340.0 + + def test_reset_budget_for_key_honors_injected_reset_time(mock_prisma_client, mock_proxy_logging): """Injected BudgetResetSettings drives the written reset time end to end (DI, no globals). diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index c547d06904b..7be6f809c0d 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1658,6 +1658,58 @@ async def test_commit_key_spend_updates_includes_last_active(): assert before_call <= last_active <= after_call +@pytest.mark.asyncio +async def test_commit_spend_updates_to_db_increments_key_total_spend_alongside_spend(): + """ + The key table write must increment the lifetime total_spend by the same amount as the + resettable spend, in the same update so the two cannot drift. + """ + db_writer = DBSpendUpdateWriter() + + mock_batcher = MagicMock() + mock_batcher.litellm_verificationtoken = MagicMock() + mock_batcher.litellm_verificationtoken.update_many = MagicMock() + + mock_transaction = AsyncMock() + mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction) + mock_transaction.__aexit__ = AsyncMock(return_value=False) + mock_transaction.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction) + + db_spend_update_transactions = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {"hashed_token_abc": 0.05, "hashed_token_def": 1.25}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + + with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=MagicMock(), + db_spend_update_transactions=db_spend_update_transactions, + ) + + calls = mock_batcher.litellm_verificationtoken.update_many.call_args_list + assert [c.kwargs["where"] for c in calls] == [{"token": "hashed_token_abc"}, {"token": "hashed_token_def"}] + for call, expected_cost in zip(calls, (0.05, 1.25)): + assert call.kwargs["data"]["spend"] == {"increment": expected_cost} + assert call.kwargs["data"]["total_spend"] == call.kwargs["data"]["spend"] + + @pytest.mark.asyncio async def test_update_database_creates_single_task(): """ @@ -2813,7 +2865,7 @@ async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at mock_batcher.litellm_verificationtoken.update_many.assert_called_once() call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1] assert call_kwargs["where"] == {"token": token} - assert set(call_kwargs["data"]) == {"spend", "last_active"} + assert set(call_kwargs["data"]) == {"spend", "total_spend", "last_active"} assert call_kwargs["data"]["spend"] == {"increment": response_cost} diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 4e70063015d..60224960bb3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1431,6 +1431,60 @@ async def test_key_info_returns_object_permission(monkeypatch): ) +def _stored_key_with_lifetime_spend(token: str, spend: float, total_spend: float) -> LiteLLM_VerificationToken: + return LiteLLM_VerificationToken.model_validate( + {"token": token, "user_id": "user123", "spend": spend, "total_spend": total_spend} + ) + + +@pytest.mark.asyncio +async def test_key_info_returns_lifetime_total_spend_next_to_resettable_spend(monkeypatch): + """After a budget reset the period spend is 0 while total_spend keeps the lifetime figure.""" + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=_stored_key_with_lifetime_spend(token="hashed_key", spend=0.0, total_spend=3.75) + ) + + result = await info_key_fn( + key="sk-test-key-456", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test-key-456"), + ) + + assert result["info"]["spend"] == 0.0 + assert result["info"]["total_spend"] == 3.75 + + +@pytest.mark.asyncio +async def test_list_keys_full_object_returns_lifetime_total_spend(): + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[_stored_key_with_lifetime_spend(token="hashed_key", spend=0.0, total_spend=3.75)] + ) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=1) + + result = await _list_key_helper( + prisma_client=mock_prisma_client, + page=1, + size=50, + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + return_full_object=True, + admin_team_ids=None, + ) + + listed_key = result["keys"][0] + assert isinstance(listed_key, UserAPIKeyAuth) + assert listed_key.spend == 0.0 + assert listed_key.total_spend == 3.75 + + @pytest.mark.asyncio async def test_get_new_token_with_valid_key(monkeypatch): """Test get_new_token function when provided with a valid key that starts with 'sk-'""" diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx index 31cd407a5e6..cf0c13ee152 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx @@ -22,6 +22,7 @@ const mockDeletedKey: DeletedKeyResponse = { key_name: "test-key", key_alias: "Test Key Alias", spend: 5.5, + total_spend: 5.5, max_budget: 100, expires: "2024-12-31T23:59:59Z", models: ["gpt-3.5-turbo"], diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 617b9209a41..8f1b7acaac0 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -79,6 +79,7 @@ const mockKey: KeyResponse = { key_name: "test-key", key_alias: "Test Key Alias", spend: 5.5, + total_spend: 42.25, max_budget: 100, expires: "2999-12-31T23:59:59Z", models: ["gpt-3.5-turbo", "gpt-4"], @@ -236,6 +237,14 @@ it("should display key information correctly", async () => { }); }); +it("shows lifetime spend in its own column next to the period spend meter", async () => { + renderWithProviders(); + + expect(await screen.findByText("Lifetime Spend")).toBeInTheDocument(); + expect(screen.getByText("$42.2500")).toBeInTheDocument(); + expect(screen.getByText("$5.5000")).toBeInTheDocument(); +}); + it("should display user email correctly", async () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index 6eea77ae827..cf5585e3486 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -13,6 +13,7 @@ import { IdCell, IdentityCell, ModelsCell, + MoneyCell, SpendBudgetCell, StatusBadge, UserPopoverCell, @@ -274,6 +275,20 @@ export const getKeyTableColumns = ({ ); }, }, + { + id: "total_spend", + accessorKey: "total_spend", + meta: { title: "Lifetime Spend" }, + header: () => ( + + ), + size: 130, + enableSorting: false, + cell: (info) => , + }, { id: "budget_reset_at", accessorKey: "budget_reset_at", diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index eadbca87140..60439e5c52c 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -39,6 +39,7 @@ export interface KeyResponse { key_name: string; key_alias: string; spend: number; + total_spend: number; max_budget: number; expires: string; models: string[]; diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index 9efcff04832..6a2f778d4a9 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -174,6 +174,7 @@ describe("KeyEditView", () => { key_name: "sk-...TUuw", key_alias: "asdasdas", spend: 0, + total_spend: 0, max_budget: 0, expires: "null", models: [], diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index b403255b329..4bf41c1f3a8 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -119,6 +119,7 @@ describe("KeyInfoView", () => { key_name: "sk-...TUuw", key_alias: "asdasdas", spend: 0, + total_spend: 0, max_budget: 0, expires: "null", models: [], @@ -272,6 +273,23 @@ describe("KeyInfoView", () => { }); }); + it("shows lifetime spend separately from the resettable period spend", async () => { + vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); + + renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + + expect(await screen.findByText("$0.2500")).toBeInTheDocument(); + expect(screen.getByTestId("key-lifetime-spend")).toHaveTextContent("Lifetime spend: $340.5000"); + }); + it("should render the key's saved router fallbacks", async () => { vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 0e6dba64110..06e088f956b 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -677,6 +677,9 @@ export default function KeyInfoView({ {currentKeyData.budget_reset_at && (

Resets {formatTimestamp(currentKeyData.budget_reset_at)}

)} +

+ Lifetime spend: ${formatNumberWithCommas(currentKeyData.total_spend ?? 0, 4)} +

@@ -935,6 +938,11 @@ export default function KeyInfoView({

${formatNumberWithCommas(currentKeyData.spend, 4)} USD

+
+

Lifetime Spend

+

${formatNumberWithCommas(currentKeyData.total_spend ?? 0, 4)} USD

+
+

Budget

diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 17ec8367324..57b0f29e3de 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29686,6 +29686,11 @@ export interface components { team_id?: string | null; /** Token */ token?: string | null; + /** + * Total Spend + * @default 0 + */ + total_spend: number; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -31259,6 +31264,11 @@ export interface components { team_id?: string | null; /** Token */ token?: string | null; + /** + * Total Spend + * @default 0 + */ + total_spend: number; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -40057,6 +40067,11 @@ export interface components { team_tpm_limit?: number | null; /** Token */ token?: string | null; + /** + * Total Spend + * @default 0 + */ + total_spend: number; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ From aebfcf7da3b5b13591232f4b559978d4f92aafb5 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 02:45:06 -0700 Subject: [PATCH 073/187] fix(e2e): route Bedrock deployments whose region only the proxy can resolve Almost every Bedrock deployment in the suite declares aws_region_name="os.environ/AWS_REGION". The mount resolver treated that string as a region name, produced a mount nothing serves, and left the whole Anthropic-on-Bedrock surface on its direct path, which is the one thing mounting Bedrock was for. The run pod does not share the proxy's environment, so the harness genuinely cannot resolve that reference. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in. --- .../test_provider_cache.py | 27 +++++++++++++------ tests/e2e/PROVIDER_CACHE.md | 2 ++ tests/e2e/provider_cache_routing.py | 23 +++++++++++++++- 3 files changed, 43 insertions(+), 9 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 5c491e1cdcd..0278bf8fd48 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -766,13 +766,20 @@ def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMPa assert route_cache_model(params, unexpected_edge, enabled=True) is params -@pytest.mark.parametrize("model", [ - "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/converse/us.anthropic.claude-sonnet-5", - "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", +@pytest.mark.parametrize("model,region", [ + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", None), + ("bedrock/converse/us.anthropic.claude-sonnet-5", None), + ("bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", None), + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "us-east-1"), + ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "os.environ/AWS_REGION"), + ("bedrock/invoke/us.anthropic.claude-sonnet-5", "os.environ/AWS_REGION"), ]) -def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str) -> None: - params: Final = LiteLLMParamsBody(model=model) +def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str, region: str | None) -> None: + """Almost every Bedrock deployment in the suite declares its region as + `os.environ/AWS_REGION`, which only the proxy can resolve. Treating that + string as a region name would leave the whole Anthropic-on-Bedrock surface + off the edge, which is the point of mounting it at all.""" + params: Final = LiteLLMParamsBody(model=model, aws_region_name=region) routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) assert routed.aws_bedrock_runtime_endpoint == "http://edge.invalid/bedrock/us-east-1" assert routed.api_base is None @@ -794,15 +801,19 @@ def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: aws_bedrock_runtime_endpoint="https://custom.invalid", ), LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_region_name="eu-west-1"), + LiteLLMParamsBody(model="bedrock/anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), + LiteLLMParamsBody(model="bedrock/invoke/eu.anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), ]) def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(params: LiteLLMParamsBody) -> None: """Non-Anthropic models the runner role cannot invoke, deployments carrying their own AWS identity (routing those would replace the assume-role chain the - batch suite exists to prove), explicit endpoints, and unmounted regions.""" + batch suite exists to prove), explicit endpoints, unmounted regions, and a + region only the proxy can resolve on a model that is not cross-region, whose + real region the harness cannot know.""" routed: Final = route_cache_model( params, lambda mount: None if mount not in EDGE_MOUNTS else f"http://edge.invalid/{mount}", enabled=True, ) - assert routed is params or routed.aws_bedrock_runtime_endpoint == params.aws_bedrock_runtime_endpoint + assert routed is params @pytest.mark.parametrize("mode", ["batch", "realtime", "image_generation"]) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 393a96e8c16..c530574cda2 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -20,6 +20,8 @@ An eligible miss calls the provider. A complete successful response is stored im Bedrock could not be mounted before because SigV4 signs the `Host` header, so a rewritten `api_base` failed signature verification at the provider. The edge now re-signs: it drops the proxy's signature headers, signs the upstream request with the run pod's own AWS identity from its EKS Pod Identity association, and forwards that. The signature headers are excluded from the key, since `x-amz-date` is a timestamp and keying on it would make every Bedrock call a permanent miss +Almost every Bedrock deployment in the suite declares its region as `os.environ/AWS_REGION`, which only the proxy can resolve, and the run pod does not share that environment. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in. + Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove. Only Anthropic models route, matching what the runner role is allowed to invoke and what the edge knows how to validate Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index e7e4899eb71..d2237bb49bc 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -10,6 +10,26 @@ LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_requ DEFAULT_BEDROCK_REGION: Final = "us-east-1" BEDROCK_ANTHROPIC_INFIX: Final = "anthropic." +BEDROCK_CROSS_REGION_PREFIX: Final = "us." +ENV_REFERENCE_PREFIX: Final = "os.environ/" + + +def bedrock_region(declared: str | None, model: str) -> str | None: + """The region whose edge mount a deployment belongs to, or None when the + harness cannot know it. + + Most Bedrock deployments declare `os.environ/AWS_REGION`, which the proxy + resolves from its own environment. The run pod does not share that + environment, so the harness genuinely does not know the region. A `us.` + inference profile fans out across the US regions and is reachable from any + of them, so the default entry point is correct for those whatever the proxy + resolved; anything else keeps its direct path rather than being sent to a + region the model may not exist in.""" + if declared is None: + return DEFAULT_BEDROCK_REGION + if not declared.startswith(ENV_REFERENCE_PREFIX): + return declared + return DEFAULT_BEDROCK_REGION if model.startswith(BEDROCK_CROSS_REGION_PREFIX) else None def bedrock_mount(params: LiteLLMParamsBody) -> str | None: @@ -24,7 +44,8 @@ def bedrock_mount(params: LiteLLMParamsBody) -> str | None: model: Final = route.partition("/")[2] or route if BEDROCK_ANTHROPIC_INFIX not in model: return None - return f"bedrock/{params.aws_region_name or DEFAULT_BEDROCK_REGION}" + region: Final = bedrock_region(params.aws_region_name, model) + return None if region is None else f"bedrock/{region}" def route_bedrock( From 6b7cafe92b7d94e9f80cb16b4d3cf3fe359801c5 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 09:58:36 +0000 Subject: [PATCH 074/187] refactor(proxy): share one typed increment for key spend and total_spend writes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_spend_update_writer.py | 11 +++++++++-- .../proxy/db/test_db_spend_update_writer.py | 13 ++++++------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 5b43ed53117..599ae90bcae 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -18,6 +18,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload from urllib.parse import quote, unquote +from typing_extensions import ReadOnly, TypedDict + import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache @@ -109,6 +111,10 @@ def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS}) +class _SpendIncrement(TypedDict): + increment: ReadOnly[float] + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -1615,11 +1621,12 @@ class DBSpendUpdateWriter: async with transaction.batch_() as batcher: # Sort by token for consistent lock ordering across pods to prevent deadlocks. for token, response_cost in sorted(key_list_transactions.items()): + spend_increment: _SpendIncrement = {"increment": response_cost} batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists where={"token": token}, data={ - "spend": {"increment": response_cost}, - "total_spend": {"increment": response_cost}, + "spend": spend_increment, + "total_spend": spend_increment, "last_active": datetime.now(timezone.utc), }, ) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 7be6f809c0d..8f72e1d6248 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1695,13 +1695,12 @@ async def test_commit_spend_updates_to_db_increments_key_total_spend_alongside_s "agent_list_transactions": {}, } - with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): - await db_writer._commit_spend_updates_to_db( - prisma_client=mock_prisma_client, - n_retry_times=0, - proxy_logging_obj=MagicMock(), - db_spend_update_transactions=db_spend_update_transactions, - ) + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=MagicMock(), + db_spend_update_transactions=db_spend_update_transactions, + ) calls = mock_batcher.litellm_verificationtoken.update_many.call_args_list assert [c.kwargs["where"] for c in calls] == [{"token": "hashed_token_abc"}, {"token": "hashed_token_def"}] From 30c6241e3a5f534037dc57a6ea544d5ff9d8bdeb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 03:01:07 -0700 Subject: [PATCH 075/187] fix(e2e): a null error field is not an error Every OpenAI Responses body carries `error: null` at the top level, and the completeness check tested the key's presence rather than its value, so it rejected every single one. The cost was silent: nothing failed, the endpoint simply never cached, which is exactly the outcome the endpoint was added for. Found by driving the edge against the real providers rather than the synthetic fixtures, which carried no error key at all. Reading the value instead of the key is also more accurate for chat completions and messages, where a real error body carries a populated error object. --- .../test_provider_cache.py | 42 +++++++++++++++++-- tests/e2e/provider_cache.py | 7 +++- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 0278bf8fd48..3dafbadf508 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -508,10 +508,13 @@ EMBEDDING_SUCCESS: Final = ( b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}],' b'"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' ) -RESPONSE_SUCCESS: Final = b'{"id":"resp_synthetic","object":"response","status":"completed","output":[]}' +RESPONSE_SUCCESS: Final = ( + b'{"id":"resp_synthetic","object":"response","status":"completed","error":null,' + b'"incomplete_details":null,"output":[]}' +) RESPONSE_STREAM_SUCCESS: Final = ( - b'data: {"type":"response.created","response":{"id":"resp_synthetic"}}\n\n' - b'data: {"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"}}\n\n' + b'data: {"type":"response.created","response":{"id":"resp_synthetic","error":null}}\n\n' + b'data: {"type":"response.completed","response":{"id":"resp_synthetic","status":"completed"},"error":null}\n\n' ) @@ -586,6 +589,39 @@ class TestNonChatOpenAiEndpoints: assert call(url, MARKED).body == payload assert len(provider.hits) == 2 + @pytest.mark.parametrize("path,response", [ + ("/v1/chat/completions", b'{"id":"x","error":null,"choices":[{"message":{"content":"hi"},' + b'"finish_reason":"stop"}]}'), + ("/v1/messages", b'{"id":"msg_x","type":"message","role":"assistant","error":null,' + b'"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}'), + ("/v1/responses", RESPONSE_SUCCESS), + ]) + def test_a_null_error_field_is_not_an_error( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + """Every OpenAI Responses body carries `error: null`, and testing the key's + presence rather than its value rejected all of them. The cost was silent: + nothing failed, the endpoint simply never cached.""" + assert b'"error":null' in response + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("path,response", [ + ("/v1/chat/completions", b'{"error":{"message":"rate limited","type":"rate_limit_error"}}'), + ("/v1/responses", b'{"object":"response","status":"completed","error":{"message":"bad"},"output":[]}'), + ]) + def test_a_populated_error_field_still_rejects( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with openai_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + @pytest.mark.parametrize("path,cacheable", [ ("/v1/chat/completions", True), ("/v1/messages", True), ("/v1/embeddings", True), ("/v1/responses", True), diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index a2b18a3a466..0dee33f33c6 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -167,7 +167,10 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, values: Final = tuple(JSON_VALUE.validate_json(event) for event in events if event != "[DONE]") except (UnicodeDecodeError, ValidationError): return False - if not values or any(not isinstance(value, dict) or "error" in value or value.get("type") == "error" for value in values): + if not values or any( + not isinstance(value, dict) or value.get("error") is not None or value.get("type") == "error" + for value in values + ): return False if urlsplit(url).path == "/v1/responses": return complete_responses_stream(values) @@ -187,7 +190,7 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, value: Final = JSON_VALUE.validate_json(body) except ValidationError: return False - if not isinstance(value, dict) or "error" in value: + if not isinstance(value, dict) or value.get("error") is not None: return False path: Final = urlsplit(url).path if path == "/v1/messages": From 4a70bc3ba366b550fdf096ffa4ccadc008ad65ba Mon Sep 17 00:00:00 2001 From: runjivu Date: Wed, 16 Sep 2026 15:19:46 +0900 Subject: [PATCH 076/187] fix: re-check budget on router fallback targets Budget is enforced once during auth, against the requested model group. `_is_model_cost_zero` waives every budget check for a zero-cost group, and the router then picks a fallback target afterwards, inside `run_async_fallback`, where nothing re-checks budget. A free model with a paid fallback therefore bills with no budget gate at all. Add `fallback_budget_check`, the budget sibling of the existing `fallback_access_check`: a predicate awaited per fallback target that skips targets the caller cannot pay for. The primary attempt is untouched, so a zero-cost model is never blocked by budget and only the paid fallback is refused. Counter reads pass `max_budget` so `get_current_spend` verifies against authoritative recorded spend, matching the auth-time key and user checks; a counter restored from an older snapshot reads as a hit rather than a clean miss, so without it a stale-low value would keep admitting paid fallbacks. A zero-cost fallback target is always allowed, and a team key does not inherit the key owner's personal budget unless `apply_user_budget_to_team_keys` is set, matching `_PROXY_MaxBudgetLimiter`. Scope is key and user budgets. Team, team-member, end-user, org, global and per-model budgets are not covered yet: those auth-path functions enforce rather than report, so reusing them would fire threshold alerts and take spend reservations for a target that is then skipped. Two limitations of that scope are documented in the module docstring: the check reads the spend counter rather than reserving against it, so concurrent fallbacks can cross a cap together; and a request reaching the router without `metadata["user_api_key_auth"]` is not restricted. Both are shared with `fallback_model_access.py`. Opt-in via `general_settings.enforce_fallback_budget`. Relates to #41344 Co-Authored-By: Claude Opus 5 (1M context) --- litellm/constants.py | 1 + litellm/proxy/auth/fallback_budget.py | 166 ++++++++++++++++ litellm/proxy/proxy_server.py | 3 + litellm/router.py | 4 + .../router_utils/fallback_event_handlers.py | 21 ++ litellm/types/router.py | 13 ++ .../proxy/auth/test_fallback_budget.py | 184 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 38 ++++ .../test_fallback_event_handlers.py | 8 + 9 files changed, 438 insertions(+) create mode 100644 litellm/proxy/auth/fallback_budget.py create mode 100644 tests/test_litellm/proxy/auth/test_fallback_budget.py diff --git a/litellm/constants.py b/litellm/constants.py index 745a4d9294e..b1fec17b6ed 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -40,6 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( "router_general_settings", "ignore_invalid_deployments", "fallback_access_check", + "fallback_budget_check", "auto_router_capability_limit", } ) diff --git a/litellm/proxy/auth/fallback_budget.py b/litellm/proxy/auth/fallback_budget.py new file mode 100644 index 00000000000..00e885d8d88 --- /dev/null +++ b/litellm/proxy/auth/fallback_budget.py @@ -0,0 +1,166 @@ +""" +Enforce the caller's budget against router fallback targets. + +Budget is checked once, during auth, against the *requested* model group. A zero-cost group takes +`_is_model_cost_zero`'s bypass and waives every budget check; the router then picks a fallback +target after auth, inside `run_async_fallback`, and nothing re-checks budget on the group that +actually bills. So a free model with a paid fallback spends without a gate. + +This predicate is injected into the router to re-check budget for each fallback target before it is +attempted, mirroring `fallback_model_access.py`. It deliberately leaves the primary attempt alone: +a zero-cost model is never blocked by budget, and only the paid fallback is refused. Opt-in via +`general_settings.enforce_fallback_budget: true`. + +Scope: the key's and the user's `max_budget`. Not covered yet, and each needs a read-only evaluation +path before it can be: team, team-member, end-user, org, global and per-model budgets, whose +auth-path functions enforce rather than report (they raise), so reusing them would fire threshold +alerts and take spend reservations for a target that is then skipped; and the key's rolling +`budget_limits` windows, whose accumulated spend lives only in per-window counters +(`spend:key:{token}:window:{budget_duration}`), so enforcing them means more counter reads on the +fallback path rather than reusing state auth already loaded. + +Two known limitations of that narrow scope, both shared with `fallback_model_access.py`: + +* This reads the spend counter, it does not reserve against it. Requests already in flight all + observe the same pre-billing figure, so a cap can be crossed by roughly the number of concurrent + fallbacks times their cost. Auth-time enforcement avoids this by pre-filling the counter through + `reserve_budget_for_request`, which the zero-cost bypass skips. Turning the soft cap into a hard + one means reserving per fallback attempt and reconciling on completion. +* A request that reaches the router without `metadata["user_api_key_auth"]` is not restricted. + Only `add_litellm_data_to_request` populates that key, so endpoints that assemble metadata by + hand (for example `/queue/chat/completions`) fall through as unauthenticated. +""" + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +from pydantic import BaseModel, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import ( + _is_model_cost_zero, # pyright: ignore[reportPrivateUsage] # the zero-cost predicate the auth-time budget checks use; no public equivalent +) +from litellm.router import Router + + +class _RequestMetadata(BaseModel): + user_api_key_auth: UserAPIKeyAuth | None = None + + +class _FallbackBudgetSettings(BaseModel): + enforce_fallback_budget: bool = False + + +def _token_in_metadata(metadata: object) -> UserAPIKeyAuth | None: + try: + return _RequestMetadata.model_validate(metadata).user_api_key_auth + except ValidationError: + return None + + +def _user_api_key_auth_from_request(request_kwargs: Mapping[str, object]) -> UserAPIKeyAuth | None: + return next( + ( + token + for field in ("metadata", "litellm_metadata") + if (token := _token_in_metadata(request_kwargs.get(field))) is not None + ), + None, + ) + + +def _enforced_by_general_settings() -> bool: + from litellm.proxy.proxy_server import general_settings + + return _FallbackBudgetSettings.model_validate(general_settings).enforce_fallback_budget + + +def _applies_user_budget_to_team_keys() -> bool: + from litellm.proxy.proxy_server import general_settings + + return general_settings.get("apply_user_budget_to_team_keys") is True + + +async def _counter_spend(counter_key: str, fallback_spend: float, max_budget: float) -> float: + """ + Read a spend counter the same way the auth-time budget checks do. + + `max_budget` is not advisory: it makes `get_current_spend` re-check the counter against the + authoritative recorded spend before admitting. A counter restored from an older Redis snapshot + reads as a hit rather than a clean miss, so without this the reseed path never runs and a + stale-low counter would keep admitting paid fallbacks past the cap. + """ + from litellm.proxy.proxy_server import get_current_spend + + return await get_current_spend( + counter_key=counter_key, + fallback_spend=fallback_spend, + max_budget=max_budget, + ) + + +async def is_token_within_budget_for_model(*, model: str, valid_token: UserAPIKeyAuth, llm_router: Router) -> bool: + """ + True when the key and the user behind it can still pay for `model`. + + A zero-cost fallback target is always allowed: refusing it would deny a request on spend some + other model accrued, which is the same reasoning behind the auth-time bypass. + """ + if _is_model_cost_zero(model=model, llm_router=llm_router): + return True + + key_budget: Final = valid_token.max_budget + if key_budget is not None and valid_token.token is not None: + key_spend: Final = await _counter_spend( + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, + max_budget=key_budget, + ) + if key_spend >= key_budget: + return False + + # Mirrors `_PROXY_MaxBudgetLimiter`: a team key does not carry the key owner's personal budget + # unless the proxy opts in, so the personal cap must not gate the fallback either. + user_budget: Final = valid_token.user_max_budget + if ( + user_budget is not None + and valid_token.user_id is not None + and (valid_token.team_id is None or _applies_user_budget_to_team_keys()) + ): + user_spend: Final = await _counter_spend( + counter_key=f"spend:user:{valid_token.user_id}", + fallback_spend=valid_token.user_spend or 0.0, + max_budget=user_budget, + ) + if user_spend >= user_budget: + return False + + return True + + +@dataclass(frozen=True, slots=True) +class RouterFallbackBudgetCheck: + """ + `FallbackBudgetCheck` for the proxy's router: while `is_enforced()` is true, a paid fallback + target is attempted only when the caller is still within budget. Requests that carry no key + (for example internal health checks) are not restricted. + """ + + is_enforced: Callable[[], bool] + + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool: + if not self.is_enforced(): + return True + valid_token: Final = _user_api_key_auth_from_request(request_kwargs) + if valid_token is None: + return True + try: + return await is_token_within_budget_for_model(model=model, valid_token=valid_token, llm_router=llm_router) + except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the caller + verbose_proxy_logger.warning("Skipping fallback to model=%s: budget lookup failed: %s", model, e) + return False + + +router_fallback_budget_check: Final = RouterFallbackBudgetCheck(is_enforced=_enforced_by_general_settings) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7964556531..0f6d84c33e7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -322,6 +322,7 @@ from litellm.proxy.auth.auth_utils import ( log_once_if_budget_reservation_disabled, warn_once_if_custom_auth_skips_common_checks, ) +from litellm.proxy.auth.fallback_budget import router_fallback_budget_check from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck @@ -6153,6 +6154,7 @@ class ProxyConfig: ), ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid fallback_access_check=router_fallback_access_check, + fallback_budget_check=router_fallback_budget_check, auto_router_capability_limit=_license_check.auto_router_capability_limit, ) @@ -6614,6 +6616,7 @@ class ProxyConfig: search_tools=search_tools, ignore_invalid_deployments=True, fallback_access_check=router_fallback_access_check, + fallback_budget_check=router_fallback_budget_check, auto_router_capability_limit=_license_check.auto_router_capability_limit, ) verbose_proxy_logger.debug("updated llm_router: %s", llm_router) diff --git a/litellm/router.py b/litellm/router.py index d531072530b..9fcda7d2e97 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -240,6 +240,7 @@ from litellm.types.router import ( DeploymentModelListingInfo, DeploymentTypedDict, FallbackAccessCheck, + FallbackBudgetCheck, GuardrailTypedDict, LiteLLM_Params, MockRouterTestingParams, @@ -755,6 +756,7 @@ class Router: background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, fallback_access_check: FallbackAccessCheck | None = None, + fallback_budget_check: FallbackBudgetCheck | None = None, auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, ) -> None: """ @@ -793,6 +795,7 @@ class Router: ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted). + fallback_budget_check (Optional[FallbackBudgetCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over budget is skipped. Defaults to None (budget is not re-checked on fallback). Returns: Router: An instance of the litellm.Router class. @@ -834,6 +837,7 @@ class Router: self.ignore_invalid_deployments = ignore_invalid_deployments self.auto_router_capability_limit = auto_router_capability_limit self.fallback_access_check: Final = fallback_access_check + self.fallback_budget_check: Final = fallback_budget_check self.debug_level = debug_level self.enable_pre_call_checks = enable_pre_call_checks self.enable_tag_filtering = enable_tag_filtering diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 94164d0ea0c..d0abaed4d3a 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -421,6 +421,25 @@ async def _is_fallback_target_authorized( return False +async def _is_fallback_target_within_budget( + litellm_router: LitellmRouter, + fallback_entry: str | Mapping[str, object], + original_model_group: str, + kwargs: Mapping[str, object], +) -> bool: + budget_check: Final = litellm_router.fallback_budget_check + target: Final = _get_fallback_target_model_group(fallback_entry) + if budget_check is None or target is None or target == original_model_group: + return True + if await budget_check(model=target, request_kwargs=kwargs, llm_router=litellm_router): + return True + verbose_router_logger.info( + "Skipping fallback to model_group = %s: caller is over budget", + mask_sensitive_structure(fallback_entry), + ) + return False + + def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: """ True when a file, batch, or fine-tuning job operation names an id that only exists @@ -528,6 +547,8 @@ async def run_async_fallback( continue if not await _is_fallback_target_authorized(litellm_router, mg, original_model_group, kwargs): continue + if not await _is_fallback_target_within_budget(litellm_router, mg, original_model_group, kwargs): + continue attempt_key = fallback_attempt_key(mg) if attempt_key is not None: if attempt_key in attempted: diff --git a/litellm/types/router.py b/litellm/types/router.py index 7c3e4d6943f..584d2494db4 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -963,6 +963,19 @@ class FallbackAccessCheck(Protocol): async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... +class FallbackBudgetCheck(Protocol): + """ + Decides whether the caller behind `request_kwargs` is still within budget for fallback `model`. + + Budget is enforced once during auth, against the *requested* model group. A fallback target is + chosen later, inside the router, so a zero-cost group that falls back to a priced one bills + without any budget gate. The router runs this before every cross-model-group fallback attempt + and skips targets it rejects, leaving the free attempt itself untouched. + """ + + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... + + class AutoRouterCapabilityLimit(Protocol): """ Resolves how many complexity routers may claim each licensed capability right now; None means unlimited. diff --git a/tests/test_litellm/proxy/auth/test_fallback_budget.py b/tests/test_litellm/proxy/auth/test_fallback_budget.py new file mode 100644 index 00000000000..0ff1e05826d --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_fallback_budget.py @@ -0,0 +1,184 @@ +import pytest + +from litellm import Router +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.fallback_budget import ( + RouterFallbackBudgetCheck, + is_token_within_budget_for_model, +) + +FREE_MODEL = { + "model_name": "free-model", + "litellm_params": { + "model": "ollama/llama2", + "api_base": "http://localhost:11434", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": { + "id": "free-model-id", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, +} + +PAID_MODEL = { + "model_name": "paid-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}, + "model_info": {"id": "paid-model-id"}, +} + + +def _router() -> Router: + return Router(model_list=[FREE_MODEL, PAID_MODEL], fallbacks=[{"free-model": ["paid-model"]}]) + + +def _token(**overrides) -> UserAPIKeyAuth: + fields = { + "api_key": "hashed", + "token": "hashed", + "spend": 0.0, + "max_budget": None, + "user_id": "u1", + "user_spend": 0.0, + "user_max_budget": None, + } + fields.update(overrides) + return UserAPIKeyAuth(**fields) + + +ENFORCED = RouterFallbackBudgetCheck(is_enforced=lambda: True) +NOT_ENFORCED = RouterFallbackBudgetCheck(is_enforced=lambda: False) + + +@pytest.mark.asyncio +async def test_paid_target_allowed_when_under_budget(): + token = _token(spend=1.0, max_budget=50.0, user_spend=1.0, user_max_budget=50.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_paid_target_refused_when_over_key_budget(): + token = _token(spend=100.0, max_budget=50.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False + + +@pytest.mark.asyncio +async def test_paid_target_refused_when_over_user_budget(): + token = _token(user_spend=1900.0, user_max_budget=50.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False + + +@pytest.mark.asyncio +async def test_zero_cost_target_allowed_even_when_over_budget(): + """Refusing a free target would deny a request on spend some other model accrued.""" + token = _token(user_spend=1900.0, user_max_budget=50.0) + assert await is_token_within_budget_for_model(model="free-model", valid_token=token, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_no_budget_configured_is_always_within_budget(): + token = _token(spend=9999.0, user_spend=9999.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_team_key_does_not_inherit_personal_budget_by_default(monkeypatch): + """Mirrors _PROXY_MaxBudgetLimiter: a team key ignores the owner's personal cap.""" + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) + token = _token(team_id="t1", user_spend=1900.0, user_max_budget=50.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_team_key_inherits_personal_budget_when_opted_in(monkeypatch): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"apply_user_budget_to_team_keys": True}, raising=False) + token = _token(team_id="t1", user_spend=1900.0, user_max_budget=50.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False + + +@pytest.mark.asyncio +async def test_check_is_a_no_op_while_not_enforced(): + request = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + assert await NOT_ENFORCED(model="paid-model", request_kwargs=request, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_request_without_a_key_is_unrestricted(): + assert await ENFORCED(model="paid-model", request_kwargs={}, llm_router=_router()) is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_enforced_check_reads_the_key_from_request_metadata(metadata_field: str): + over = {metadata_field: {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + under = {metadata_field: {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}} + + assert await ENFORCED(model="paid-model", request_kwargs=over, llm_router=_router()) is False + assert await ENFORCED(model="paid-model", request_kwargs=under, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_a_stale_low_counter_still_refuses_a_paid_target(monkeypatch): + """ + The counter can read low (e.g. restored from an older Redis snapshot). Passing the budget makes + `get_current_spend` verify against authoritative spend instead of trusting that read, so the + paid target is still refused. + """ + from litellm.proxy import proxy_server + + seen: list[dict] = [] + + async def _stale_counter(**kwargs): + seen.append(kwargs) + # a stale-low counter read; the authoritative spend is what the budget must be judged on + return 0.0 if kwargs.get("max_budget") is None else kwargs["fallback_spend"] + + monkeypatch.setattr(proxy_server, "get_current_spend", _stale_counter, raising=False) + token = _token(user_spend=1900.0, user_max_budget=50.0) + + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False + assert [call["max_budget"] for call in seen] == [50.0] + + +@pytest.mark.asyncio +async def test_check_fails_closed_when_the_spend_lookup_breaks(monkeypatch): + from litellm.proxy import proxy_server + + async def _boom(**kwargs): + raise RuntimeError("spend counter unavailable") + + monkeypatch.setattr(proxy_server, "get_current_spend", _boom, raising=False) + request = {"metadata": {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}} + + assert await ENFORCED(model="paid-model", request_kwargs=request, llm_router=_router()) is False + + +@pytest.mark.asyncio +async def test_router_skips_the_paid_fallback_target_when_over_budget(): + from litellm.router_utils.fallback_event_handlers import _is_fallback_target_within_budget + + router = Router( + model_list=[FREE_MODEL, PAID_MODEL], + fallbacks=[{"free-model": ["paid-model"]}], + fallback_budget_check=ENFORCED, + ) + over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + under = {"metadata": {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}} + + assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", over) is False + assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", under) is True + + +@pytest.mark.asyncio +async def test_router_without_a_budget_check_attempts_every_fallback(): + from litellm.router_utils.fallback_event_handlers import _is_fallback_target_within_budget + + router = _router() # fallback_budget_check defaults to None + over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + + assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", over) is True diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 6f55449abab..ca53887623f 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13387,6 +13387,44 @@ async def test_load_config_router_authorizes_fallback_targets_against_the_callin assert router.fallback_access_check is router_fallback_access_check +@pytest.mark.asyncio +async def test_load_config_router_budget_checks_fallback_targets_against_the_calling_key(tmp_path, monkeypatch): + """A config-loaded router refuses a paid fallback target for an over-budget caller.""" + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import ProxyConfig + + config_file = tmp_path / "config.yaml" + config_file.write_text( + yaml.dump({"model_list": [{"model_name": "m", "litellm_params": {"model": "openai/m", "api_key": "k"}}]}) + ) + + router, _, _ = await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + over_budget = { + "metadata": { + "user_api_key_auth": UserAPIKeyAuth( + api_key="hashed", token="hashed", user_id="u1", user_spend=99.0, user_max_budget=1.0 + ) + } + } + under_budget = { + "metadata": { + "user_api_key_auth": UserAPIKeyAuth( + api_key="hashed", token="hashed", user_id="u1", user_spend=0.0, user_max_budget=100.0 + ) + } + } + + # off by default: the paid fallback is still attempted for an over-budget caller + monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) + assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, llm_router=router) is True + + monkeypatch.setattr(proxy_server, "general_settings", {"enforce_fallback_budget": True}, raising=False) + assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, llm_router=router) is False + assert await router.fallback_budget_check(model="m", request_kwargs=under_budget, llm_router=router) is True + + @pytest.mark.asyncio async def test_load_config_user_api_key_cache_max_size_keeps_more_than_200_entries(tmp_path, monkeypatch): """The auth cache used to be pinned at InMemoryCache's 200 entry default, so a diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 9318f306c89..dfe06bffd09 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -27,6 +27,7 @@ class StreamingWrapper: class FakeRouter: fallback_access_check = None + fallback_budget_check = None def log_retry(self, kwargs, e): return kwargs @@ -37,6 +38,7 @@ class FakeRouter: class AlwaysFailRouter: fallback_access_check = None + fallback_budget_check = None def log_retry(self, kwargs, e): return kwargs @@ -101,6 +103,7 @@ async def test_run_async_fallback_raises_when_all_fallbacks_fail(): class RecordingRouter: fallback_access_check = None + fallback_budget_check = None def __init__(self): self.received_kwargs = None @@ -162,6 +165,7 @@ async def test_run_async_fallback_skips_original_model_group(): class AttemptRecordingRouter: fallback_access_check = None + fallback_budget_check = None def __init__(self): self.attempted_model_groups = [] @@ -471,6 +475,8 @@ class AccessCheckedRouter(AttemptRecordingRouter): self.allowed_models = allowed_models self.access_checks = [] + fallback_budget_check = None + async def fallback_access_check(self, *, model, request_kwargs, llm_router): self.access_checks.append((model, request_kwargs["metadata"]["user_api_key"], llm_router is self)) return model in self.allowed_models @@ -542,6 +548,7 @@ async def test_run_async_fallback_does_not_consult_access_check_for_same_model_g class RecordingFailRouter: fallback_access_check = None + fallback_budget_check = None def __init__(self): self.attempted_models = [] @@ -1053,6 +1060,7 @@ class TestTriggerCooldownForFailedDeployment: class TestRunAsyncFallbackTriggersCooldown: class RouterWithLoggingKwarg: fallback_access_check = None + fallback_budget_check = None def __init__(self): self.cooldown_time = 60.0 From c7246adc1dd571069a2cb3c77b3a40b720eb0da6 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 03:29:02 -0700 Subject: [PATCH 077/187] fix(e2e): route only the Bedrock models the runner role can invoke The edge re-signs with the run pod's identity, whose IAM policy is an explicit per-model allowlist. Matching on the `anthropic.` infix instead routed every Anthropic-on-Bedrock model, so a model outside the policy came back 403 from Bedrock with no fallback, taking the whole claude_code Bedrock matrix red. An unlisted model now keeps its direct path and loses only caching. --- .../test_provider_cache.py | 4 ++++ tests/e2e/PROVIDER_CACHE.md | 4 +++- tests/e2e/provider_cache_routing.py | 23 ++++++++++++------- 3 files changed, 22 insertions(+), 9 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 3dafbadf508..ac1ccd5692b 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -809,6 +809,8 @@ def test_registration_preserves_unsupported_or_explicit_routes(params: LiteLLMPa ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "us-east-1"), ("bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", "os.environ/AWS_REGION"), ("bedrock/invoke/us.anthropic.claude-sonnet-5", "os.environ/AWS_REGION"), + ("bedrock/us.anthropic.claude-opus-4-7", "us-east-1"), + ("bedrock/converse/us.anthropic.claude-opus-4-7", "us-east-1"), ]) def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: str, region: str | None) -> None: """Almost every Bedrock deployment in the suite declares its region as @@ -839,6 +841,8 @@ def test_anthropic_on_bedrock_registers_the_edge_as_its_runtime_endpoint(model: LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", aws_region_name="eu-west-1"), LiteLLMParamsBody(model="bedrock/anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), LiteLLMParamsBody(model="bedrock/invoke/eu.anthropic.claude-sonnet-5", aws_region_name="os.environ/AWS_REGION"), + LiteLLMParamsBody(model="bedrock/us.anthropic.claude-opus-4-5", aws_region_name="us-east-1"), + LiteLLMParamsBody(model="bedrock/converse/us.anthropic.claude-haiku-9-9", aws_region_name="us-east-1"), ]) def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(params: LiteLLMParamsBody) -> None: """Non-Anthropic models the runner role cannot invoke, deployments carrying diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index c530574cda2..3a29c7b6fe4 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -22,7 +22,9 @@ Bedrock could not be mounted before because SigV4 signs the `Host` header, so a Almost every Bedrock deployment in the suite declares its region as `os.environ/AWS_REGION`, which only the proxy can resolve, and the run pod does not share that environment. A `us.` inference profile fans out across the US regions and is reachable from any of them, so those route to the default mount whatever the proxy resolved. A model that is not cross-region and declares its region that way keeps its direct path rather than being sent to a region it may not exist in. -Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove. Only Anthropic models route, matching what the runner role is allowed to invoke and what the edge knows how to validate +Only deployments that carry no AWS identity of their own route to the edge. A deployment with `aws_role_name`, `aws_access_key_id`, an `api_base` or an `aws_bedrock_runtime_endpoint` keeps its direct path, because re-signing it would quietly replace the very credential chain that test exists to prove + +Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index d2237bb49bc..97e05344423 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -9,8 +9,15 @@ from models import LiteLLMParamsBody, ModelMode LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False) DEFAULT_BEDROCK_REGION: Final = "us-east-1" -BEDROCK_ANTHROPIC_INFIX: Final = "anthropic." BEDROCK_CROSS_REGION_PREFIX: Final = "us." +BEDROCK_EDGE_MODELS: Final = frozenset( + { + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "us.anthropic.claude-sonnet-5", + "us.anthropic.claude-opus-4-7", + } +) ENV_REFERENCE_PREFIX: Final = "os.environ/" @@ -33,16 +40,16 @@ def bedrock_region(declared: str | None, model: str) -> str | None: def bedrock_mount(params: LiteLLMParamsBody) -> str | None: - """The edge mount an Anthropic-on-Bedrock deployment belongs to, or None. + """The edge mount a Bedrock deployment belongs to, or None. - Only the Anthropic models route. The edge validates converse and invoke - bodies by their Anthropic and Converse terminator fields, and the runner role - is allowed to invoke exactly those models, so Bedrock embeddings, image - generation, rerank and realtime keep their existing direct path rather than - reaching an edge that could neither sign nor validate for them.""" + The allowlist mirrors the runner role's IAM policy, which names its models + one by one. A model outside it would be re-signed with an identity that + cannot invoke it and come back 403 from Bedrock, so an unlisted model keeps + its direct path and loses only caching. Adding a model is a policy edit in + litellm-ops and a line here.""" route: Final = params.model.partition("/")[2] model: Final = route.partition("/")[2] or route - if BEDROCK_ANTHROPIC_INFIX not in model: + if model not in BEDROCK_EDGE_MODELS: return None region: Final = bedrock_region(params.aws_region_name, model) return None if region is None else f"bedrock/{region}" From ebf34cd88006110c78a50648578e3dc7e0f115cd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 03:35:00 -0700 Subject: [PATCH 078/187] docs(e2e): say plainly that Bedrock streaming is not cached --- tests/e2e/PROVIDER_CACHE.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 3a29c7b6fe4..fe289e406aa 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -2,7 +2,9 @@ `E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live -The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, including streams. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored + +Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, are not cacheable. They still cross the edge and are still re-signed, so they need the same IAM, but they always call the provider. AWS frames them as binary `vnd.amazon.eventstream` rather than SSE, and reading a terminal event out of that is what a completeness rule for them would need. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so most Bedrock traffic in the suite is not cached today ## Request identity From 7c2234be3a91a600c02f416e678b163dd53746ac Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 04:41:29 -0700 Subject: [PATCH 079/187] test(e2e): enforce the cross-region invariant on the Bedrock allowlist The allowlist rejects an unlisted model before the region resolver runs, so the two negative cases that used to cover the resolver were passing for the wrong reason and two mutations of it survived. Answering an env-referenced region with the default mount is only sound because every allowlisted model is a `us.` profile that fans out across the US regions, so assert that on the list itself and drop the per-call branch it made unreachable. --- .../test_provider_cache.py | 33 ++++++++++++++++++- tests/e2e/provider_cache_routing.py | 27 +++++++-------- 2 files changed, 43 insertions(+), 17 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index ac1ccd5692b..c24ba8d6221 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -33,7 +33,13 @@ from provider_cache import ( successful_response, ) from provider_cache_redis import PUBLISH, RedisCommands, RedisResponseStore, configured_cache, redis_store -from provider_cache_routing import LIVE_PROVIDER_REQUIRED, route_cache_model +from provider_cache_routing import ( + BEDROCK_CROSS_REGION_PREFIX, + BEDROCK_EDGE_MODELS, + LIVE_PROVIDER_REQUIRED, + bedrock_region, + route_cache_model, +) from fixture_mode import SESSION_TEST_KEY from provider_edge import EDGE_MOUNTS, configured_cache_backend, resolve_mount, start_provider_edge from provider_edge_bedrock import bedrock_signer @@ -856,6 +862,31 @@ def test_bedrock_deployments_the_edge_must_not_touch_keep_their_direct_route(par assert routed is params +@pytest.mark.parametrize("declared,expected", [ + (None, "us-east-1"), + ("us-west-2", "us-west-2"), + ("eu-west-1", "eu-west-1"), + ("os.environ/AWS_REGION", "us-east-1"), + ("os.environ/ANY_OTHER_NAME", "us-east-1"), +]) +def test_a_region_only_the_proxy_can_resolve_falls_back_to_the_default_mount( + declared: str | None, expected: str, +) -> None: + """A declared literal region is the one the deployment meant. A region the + proxy resolves from its own environment is one the run pod cannot see, and + the default mount answers it.""" + assert bedrock_region(declared) == expected + + +def test_every_model_on_the_edge_allowlist_is_a_cross_region_profile() -> None: + """Answering an env-referenced region with the default mount is only correct + for a profile that fans out across the US regions and is reachable from any + of them. A single-region model on this list would be sent to a region it may + not exist in, so the list is where that is caught.""" + assert BEDROCK_EDGE_MODELS + assert all(model.startswith(BEDROCK_CROSS_REGION_PREFIX) for model in BEDROCK_EDGE_MODELS) + + @pytest.mark.parametrize("mode", ["batch", "realtime", "image_generation"]) def test_a_bedrock_deployment_with_a_mode_keeps_its_direct_route(mode: ModelMode) -> None: params: Final = LiteLLMParamsBody(model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0") diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index 97e05344423..f9775a2b152 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -21,22 +21,18 @@ BEDROCK_EDGE_MODELS: Final = frozenset( ENV_REFERENCE_PREFIX: Final = "os.environ/" -def bedrock_region(declared: str | None, model: str) -> str | None: - """The region whose edge mount a deployment belongs to, or None when the - harness cannot know it. +def bedrock_region(declared: str | None) -> str: + """The region whose edge mount a deployment belongs to. - Most Bedrock deployments declare `os.environ/AWS_REGION`, which the proxy - resolves from its own environment. The run pod does not share that - environment, so the harness genuinely does not know the region. A `us.` - inference profile fans out across the US regions and is reachable from any - of them, so the default entry point is correct for those whatever the proxy - resolved; anything else keeps its direct path rather than being sent to a - region the model may not exist in.""" - if declared is None: + Most Bedrock deployments declare `os.environ/AWS_REGION`, which only the + proxy can resolve from its own environment; the run pod does not share it. + Answering those with the default mount is correct because every model on the + edge allowlist is a `us.` inference profile, which fans out across the US + regions and is reachable from any of them. That invariant is enforced on the + allowlist itself rather than re-checked per call.""" + if declared is None or declared.startswith(ENV_REFERENCE_PREFIX): return DEFAULT_BEDROCK_REGION - if not declared.startswith(ENV_REFERENCE_PREFIX): - return declared - return DEFAULT_BEDROCK_REGION if model.startswith(BEDROCK_CROSS_REGION_PREFIX) else None + return declared def bedrock_mount(params: LiteLLMParamsBody) -> str | None: @@ -51,8 +47,7 @@ def bedrock_mount(params: LiteLLMParamsBody) -> str | None: model: Final = route.partition("/")[2] or route if model not in BEDROCK_EDGE_MODELS: return None - region: Final = bedrock_region(params.aws_region_name, model) - return None if region is None else f"bedrock/{region}" + return f"bedrock/{bedrock_region(params.aws_region_name)}" def route_bedrock( From bd1c2d6f07d7b9edd46bb11d95ad22cb7e029ea9 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 05:09:04 -0700 Subject: [PATCH 080/187] fix(e2e): keep the tool-continuation echo-back test on the live path The key normalizes a unique marker so two builds match, which is the whole point, but it makes this test's identity collide with an earlier run's: it mints a fresh receipt, sends it through a tool result, and asserts the model echoes it back verbatim, so a stale recording matched and answered with the old receipt. Build 223 is where that surfaced, once the corpus was full enough for the first call to hit. A test that asserts a provider echoed this run's own unique value belongs on the live path. --- tests/e2e/PROVIDER_CACHE.md | 4 +++- tests/e2e/conftest.py | 6 +++++- tests/e2e/llm_translation/test_messages_e2e.py | 1 + 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index fe289e406aa..698e78a5aa1 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -44,7 +44,9 @@ The trusted runner receives: Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits -Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay +Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. + +One more class needs it, and it is the cost of normalizing the marker. A test that mints a fresh marker, sends it, and then asserts the provider's answer contains that exact value is asserting on the marker rather than using it as a salt. The key treats two such requests as the same identity, so a stale recording matches and answers with the marker from the run that recorded it. `TestOpenAIMessagesToolContinuation` is the one in the suite today: it sends a freshly minted receipt through a tool result and asserts the model echoes it back verbatim. If you add a test that asserts a provider echoed your own unique value, it belongs on the live path. Ordinary assertion failures still fail E2E. The shared cache does not modify provider response IDs or make the proxy aware of replay ## Recorded response semantics diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 829c84910a9..430e16525d5 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -85,7 +85,11 @@ def jwt_identity(idp: Keycloak, resources: ResourceManager, proxy: ProxyClient) def pytest_configure(config: pytest.Config) -> None: - config.addinivalue_line("markers", "provider_live: requires actual provider timing, limits or state; bypass shared cache") + config.addinivalue_line( + "markers", + "provider_live: requires actual provider timing, limits, state, or a response that echoes this" + " run's own unique value; bypass shared cache", + ) config.addinivalue_line( "markers", "e2e: live test that requires a running proxy and real provider keys", diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index 44c416a3e78..09ec48daa2f 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -372,6 +372,7 @@ def _request_tool( class TestOpenAIMessagesToolContinuation: + @pytest.mark.provider_live @pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"]) def test_required_tool_arguments_and_correlated_result( self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool From 642525f5712b50c9847eecf17c32d39152d50b8f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:19:53 +0000 Subject: [PATCH 081/187] fix(registry): add azure gpt-image-2.5 entries and together/azure deprecation dates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 53 +++++++++++++++++-- model_prices_and_context_window.json | 53 +++++++++++++++++-- 2 files changed, 98 insertions(+), 8 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index eafbb8b1048..542a1113112 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5265,6 +5265,7 @@ "supports_vision": true }, "azure/gpt-4o": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -8988,6 +8989,36 @@ "supports_vision": true, "supports_pdf_input": true }, + "azure/gpt-image-2.5-flare": { + "deprecation_date": "2027-09-09", + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "azure/gpt-image-2.5-sunburst": { + "deprecation_date": "2027-09-09", + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "azure/gpt-image-2-2026-04-21": { "cache_read_input_token_cost": 1.25e-06, "deprecation_date": "2027-10-21", @@ -44437,7 +44468,7 @@ "supports_tool_choice": true }, "together_ai/openai/gpt-oss-20b": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "input_cost_per_token": 5e-08, "litellm_provider": "together_ai", "max_input_tokens": 131072, @@ -44744,7 +44775,7 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/google/gemma-4-31B-it": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -44759,7 +44790,7 @@ "supports_vision": true }, "together_ai/intfloat/multilingual-e5-large-instruct": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "input_cost_per_token": 2e-08, "litellm_provider": "together_ai", "max_input_tokens": 514, @@ -44872,7 +44903,7 @@ "supports_tool_choice": true }, "together_ai/thinkingmachines/Inkling-Small": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", @@ -66337,6 +66368,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/deepseek-coder-33b-instruct": { + "deprecation_date": "2024-08-22", "input_cost_per_token": 8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66344,6 +66376,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { + "deprecation_date": "2025-12-23", "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", "mode": "chat", @@ -66351,6 +66384,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 1.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66358,6 +66392,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 1.6e-06, "litellm_provider": "together_ai", "mode": "chat", @@ -66443,6 +66478,7 @@ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, "together_ai/google/gemma-2-27b-it": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66467,6 +66503,7 @@ "source": "https://developers.openai.com/api/docs/pricing" }, "together_ai/meta-llama/Llama-3-8b-chat-hf": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66495,6 +66532,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/meta-llama/Meta-Llama-3-70B-Instruct-Turbo": { + "deprecation_date": "2025-12-23", "input_cost_per_token": 8.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66502,6 +66540,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/meta-llama/Meta-Llama-3-8B-Instruct": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66509,6 +66548,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66516,6 +66556,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/nvidia/Llama-3.1-Nemotron-70B-Instruct-HF": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 8.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66530,6 +66571,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen2-72B-Instruct": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 9e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66537,6 +66579,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen2-VL-72B-Instruct": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 1.2e-06, "litellm_provider": "together_ai", "mode": "chat", @@ -66558,6 +66601,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen2.5-Coder-32B-Instruct": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66565,6 +66609,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen2.5-VL-72B-Instruct": { + "deprecation_date": "2026-01-05", "input_cost_per_token": 1.95e-06, "litellm_provider": "together_ai", "mode": "chat", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index eafbb8b1048..542a1113112 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5265,6 +5265,7 @@ "supports_vision": true }, "azure/gpt-4o": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -8988,6 +8989,36 @@ "supports_vision": true, "supports_pdf_input": true }, + "azure/gpt-image-2.5-flare": { + "deprecation_date": "2027-09-09", + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "azure/gpt-image-2.5-sunburst": { + "deprecation_date": "2027-09-09", + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "azure/gpt-image-2-2026-04-21": { "cache_read_input_token_cost": 1.25e-06, "deprecation_date": "2027-10-21", @@ -44437,7 +44468,7 @@ "supports_tool_choice": true }, "together_ai/openai/gpt-oss-20b": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "input_cost_per_token": 5e-08, "litellm_provider": "together_ai", "max_input_tokens": 131072, @@ -44744,7 +44775,7 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/google/gemma-4-31B-it": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -44759,7 +44790,7 @@ "supports_vision": true }, "together_ai/intfloat/multilingual-e5-large-instruct": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "input_cost_per_token": 2e-08, "litellm_provider": "together_ai", "max_input_tokens": 514, @@ -44872,7 +44903,7 @@ "supports_tool_choice": true }, "together_ai/thinkingmachines/Inkling-Small": { - "deprecation_date": "2026-09-14", + "deprecation_date": "2026-09-15", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", @@ -66337,6 +66368,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/deepseek-coder-33b-instruct": { + "deprecation_date": "2024-08-22", "input_cost_per_token": 8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66344,6 +66376,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { + "deprecation_date": "2025-12-23", "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", "mode": "chat", @@ -66351,6 +66384,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 1.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66358,6 +66392,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 1.6e-06, "litellm_provider": "together_ai", "mode": "chat", @@ -66443,6 +66478,7 @@ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, "together_ai/google/gemma-2-27b-it": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66467,6 +66503,7 @@ "source": "https://developers.openai.com/api/docs/pricing" }, "together_ai/meta-llama/Llama-3-8b-chat-hf": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66495,6 +66532,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/meta-llama/Meta-Llama-3-70B-Instruct-Turbo": { + "deprecation_date": "2025-12-23", "input_cost_per_token": 8.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66502,6 +66540,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/meta-llama/Meta-Llama-3-8B-Instruct": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66509,6 +66548,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66516,6 +66556,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/nvidia/Llama-3.1-Nemotron-70B-Instruct-HF": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 8.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66530,6 +66571,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen2-72B-Instruct": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 9e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66537,6 +66579,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen2-VL-72B-Instruct": { + "deprecation_date": "2025-08-28", "input_cost_per_token": 1.2e-06, "litellm_provider": "together_ai", "mode": "chat", @@ -66558,6 +66601,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen2.5-Coder-32B-Instruct": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -66565,6 +66609,7 @@ "source": "https://api.together.ai/v1/models" }, "together_ai/Qwen/Qwen2.5-VL-72B-Instruct": { + "deprecation_date": "2026-01-05", "input_cost_per_token": 1.95e-06, "litellm_provider": "together_ai", "mode": "chat", From ba820783acc17406aa6ffa50ea53c64e68d0ed6a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 13:25:14 +0000 Subject: [PATCH 082/187] fix(tests): import Bedrock usage types Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock/test_cross_region_inference_profile_mapping.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index 65e38c148be..5795e29a8bc 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -12,8 +12,9 @@ from litellm.types.utils import ( Choices, Message, ModelResponse, + PromptTokensDetailsWrapper, + Usage, ) -from litellm.utils import _get_model_info_helper @pytest.fixture From 30a691ed55c59d8fd19c28e5356ba196bf5ae45a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 06:46:19 -0700 Subject: [PATCH 083/187] feat(e2e): cache Bedrock streaming responses The Claude Code compat cells drive the real CLI, which always streams, so converse-stream and invoke-with-response-stream were most of the suite's Bedrock traffic and all of it bypassed the edge. AWS frames those as binary vnd.amazon.eventstream rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. Two details drove the rule. A ConverseStream ends with metadata, not with messageStop, and metadata is what carries the token usage litellm prices the call from, so a stream cut between the two names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser: it yields the frames it did receive and silently discards a trailing partial one, so a stream cut one byte short parses clean. The body is checked against the frame lengths it declares to catch that. The invoke stream carries the ordinary Anthropic event grammar inside its chunk frames, so it shares the completeness rule with the SSE mounts. Validated against three real Bedrock eventstream captures, and the tests build their own frames rather than pasting a capture, with one test holding that framing to botocore's parser. --- .../test_provider_cache.py | 258 ++++++++++++++++-- tests/e2e/PROVIDER_CACHE.md | 4 +- tests/e2e/provider_cache.py | 141 +++++++++- 3 files changed, 370 insertions(+), 33 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index c24ba8d6221..4c131434ecd 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -1,9 +1,13 @@ from __future__ import annotations +import base64 +import binascii +import json import os import shutil import socket import subprocess +import struct import threading import time import uuid @@ -17,9 +21,11 @@ from typing import Final from urllib.parse import urlsplit import pytest +from pydantic import JsonValue from e2e_http import NetworkError, PreparedForward, RawResponse, StreamChunk, StreamHead, forward, prepare_forward from models import LiteLLMParamsBody, ModelMode from botocore.credentials import Credentials +from botocore.eventstream import EventStreamBuffer from provider_cache import ( SIGNATURE_HEADERS, CacheEdge, @@ -638,6 +644,53 @@ class TestNonChatOpenAiEndpoints: assert cacheable_endpoint("openai", "POST", f"https://api.openai.com{path}", MARKED) is cacheable +BEDROCK_STREAM_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1:0" +CONVERSE_STREAM_URL: Final = f"https://bedrock.invalid/model/{BEDROCK_STREAM_MODEL}/converse-stream" +INVOKE_STREAM_URL: Final = f"https://bedrock.invalid/model/{BEDROCK_STREAM_MODEL}/invoke-with-response-stream" + + +def eventstream_frame(headers: Mapping[str, str], payload: bytes) -> bytes: + """AWS eventstream wire framing, the shape `vnd.amazon.eventstream` bodies + arrive in. Built here rather than pasted from a capture so a test can express + the stream it means; `test_the_frames_these_tests_build_are_real_aws_framing` + holds it to botocore's own parser.""" + encoded: Final = b"".join( + bytes([len(name)]) + name.encode() + b"\x07" + struct.pack(">H", len(value)) + value.encode() + for name, value in headers.items() + ) + prelude: Final = struct.pack(">II", 16 + len(encoded) + len(payload), len(encoded)) + framed: Final = prelude + struct.pack(">I", binascii.crc32(prelude)) + encoded + payload + return framed + struct.pack(">I", binascii.crc32(framed)) + + +def eventstream_event(event_type: str, payload: JsonValue, message_type: str = "event") -> bytes: + return eventstream_frame( + {":event-type": event_type, ":message-type": message_type, ":content-type": "application/json"}, + json.dumps(payload).encode(), + ) + + +def invoke_chunk(inner: JsonValue) -> bytes: + return eventstream_event("chunk", {"bytes": base64.b64encode(json.dumps(inner).encode()).decode("ascii")}) + + +CONVERSE_STREAM_OK: Final = ( + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "hi"}}) + + eventstream_event("contentBlockStop", {"contentBlockIndex": 0}) + + eventstream_event("messageStop", {"stopReason": "end_turn"}) + + eventstream_event("metadata", {"usage": {"inputTokens": 12, "outputTokens": 6, "totalTokens": 18}}) +) +INVOKE_STREAM_OK: Final = ( + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x", "role": "assistant"}}) + + invoke_chunk({"type": "content_block_start", "index": 0}) + + invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"text": "hi"}}) + + invoke_chunk({"type": "content_block_stop", "index": 0}) + + invoke_chunk({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}) + + invoke_chunk({"type": "message_stop"}) +) + + class TestBedrockSigning: """Bedrock is the reason the edge could not mount it before: SigV4 covers the Host header, so forwarding through a rewritten api_base invalidates the @@ -738,32 +791,55 @@ class TestBedrockSigning: assert call(url, BEDROCK_BODY).body == response assert len(provider.hits) == 2 - @pytest.mark.parametrize("action", ["converse-stream", "invoke-with-response-stream"]) - def test_streaming_endpoints_go_live_every_time( - self, store: RedisResponseStore, provider: Provider, action: str, + @pytest.mark.parametrize("action,response", [ + ("converse-stream", CONVERSE_STREAM_OK), + ("invoke-with-response-stream", INVOKE_STREAM_OK), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_finished_stream_is_served_from_the_cache_the_second_time( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, ) -> None: - """An eventstream's completeness cannot be proven without parsing its - frames, so these bypass rather than risk recording a truncated answer. - They are still signed: a bypass is a forward, not a passthrough.""" - provider.response = CONVERSE_SUCCESS - cache: Final = bedrock_cache_edge(store) - for _ in range(2): - with bedrock_edge(cache, provider, action) as url: - assert call(url, BEDROCK_BODY).body == CONVERSE_SUCCESS - assert len(provider.hits) == 2 - assert dict(cache.counters.counts)[f"mount:{BEDROCK_MOUNT}:bypass"] == 2 + provider.response = response + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + replay: Final = bedrock_cache_edge(store) + with bedrock_edge(replay, provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 1 + assert dict(replay.counters.counts)[f"mount:{BEDROCK_MOUNT}:hits"] == 1 assert all( sent.startswith("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/") for sent in provider.authorizations ), provider.authorizations - @pytest.mark.parametrize("action,cacheable", [ - ("converse", True), ("invoke", True), - ("converse-stream", False), ("invoke-with-response-stream", False), - ]) - def test_only_the_unary_bedrock_actions_are_cacheable(self, action: str, cacheable: bool) -> None: + @pytest.mark.parametrize("action,response", [ + ("converse-stream", CONVERSE_STREAM_OK[:-1]), + ("invoke-with-response-stream", INVOKE_STREAM_OK[:-1]), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_stream_the_connection_cut_short_calls_the_provider_every_time( + self, store: RedisResponseStore, provider: Provider, action: str, response: bytes, + ) -> None: + """The whole risk of caching an eventstream is recording a half-finished + one, so a truncated body has to be rejected rather than stored.""" + provider.response = response + with bedrock_edge(bedrock_cache_edge(store), provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + replay: Final = bedrock_cache_edge(store) + with bedrock_edge(replay, provider, action) as url: + assert call(url, BEDROCK_BODY).body == response + assert len(provider.hits) == 2 + assert dict(replay.counters.counts)[f"mount:{BEDROCK_MOUNT}:rejected"] == 1 + assert f"mount:{BEDROCK_MOUNT}:hits" not in dict(replay.counters.counts) + + @pytest.mark.parametrize("action", ["converse", "invoke", "converse-stream", "invoke-with-response-stream"]) + def test_every_anthropic_bedrock_action_is_cacheable(self, action: str) -> None: url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" - assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) is cacheable + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) + + @pytest.mark.parametrize("action", ["count-tokens", "invoke-async", "converse-stream-x"]) + def test_an_unknown_bedrock_action_is_not_cacheable(self, action: str) -> None: + url: Final = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_MODEL}/{action}" + assert not cacheable_endpoint(BEDROCK_MOUNT, "POST", url, BEDROCK_BODY) def test_a_region_mount_resolves_whole(self) -> None: resolved: Final = resolve_mount(f"/{BEDROCK_MOUNT}/model/{BEDROCK_MODEL}/converse", EDGE_MOUNTS) @@ -1021,3 +1097,147 @@ def test_duplicate_headers_bypass_cache_and_count_live_calls( assert len(provider.hits) == (2 if known_mount else 0) assert dict(cache.counters.counts)["duplicate_header_bypass"] == 2 assert dict(cache.counters.counts).get("upstream_attempts", 0) == (2 if known_mount else 0) + + +class TestBedrockStreams: + def test_the_frames_these_tests_build_are_real_aws_framing(self) -> None: + buffer: Final = EventStreamBuffer() + buffer.add_data(CONVERSE_STREAM_OK) + assert [event.headers[":event-type"] for event in buffer] == [ + "messageStart", "contentBlockDelta", "contentBlockStop", "messageStop", "metadata", + ] + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + def test_a_finished_stream_is_recordable(self, url: str, body: bytes) -> None: + assert cacheable_endpoint(BEDROCK_MOUNT, "POST", url, b"{}") + assert successful_response(BEDROCK_MOUNT, url, 200, {}, body) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + @pytest.mark.parametrize("keep", [1, -1, -4]) + def test_a_stream_the_connection_cut_short_is_not_recordable( + self, url: str, body: bytes, keep: int, + ) -> None: + """botocore yields the frames it did receive and silently drops a trailing + partial one, so a stream cut a single byte short parses clean and only the + byte accounting and the terminator rule catch it.""" + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, body[:keep]) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ]) + def test_a_corrupted_frame_is_not_recordable(self, url: str, body: bytes) -> None: + flipped: Final = bytearray(body) + flipped[len(body) // 2] ^= 0xFF + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, bytes(flipped)) + + def test_a_converse_stream_that_lost_its_usage_is_not_recordable(self) -> None: + """ConverseStream names its stop reason a frame before it reports usage, + and litellm prices the call from that usage, so a stream cut between the + two would replay as a free call.""" + without_metadata: Final = ( + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("messageStop", {"stopReason": "end_turn"}) + ) + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, without_metadata) + + def test_a_converse_stream_that_never_stopped_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}), + ) + + def test_a_stream_that_failed_after_answering_200_is_not_recordable(self) -> None: + """Bedrock reports a fault that began after the headers went out as an + exception frame in place of the terminator it never got to send.""" + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "hi"}}) + + eventstream_event("modelStreamErrorException", {"message": "boom"}, message_type="exception"), + ) + + @pytest.mark.parametrize("url,body", [ + (CONVERSE_STREAM_URL, CONVERSE_STREAM_OK), + (INVOKE_STREAM_URL, INVOKE_STREAM_OK), + ], ids=["converse-stream", "invoke-stream"]) + def test_a_stream_cut_after_its_terminator_is_not_recordable(self, url: str, body: bytes) -> None: + """The terminator rules cannot see this one. Every frame the stream owes + has arrived and the partial frame after them is the one botocore drops + without a word, so only counting the bytes against the frame lengths + tells this from a stream that ended where it meant to.""" + assert successful_response(BEDROCK_MOUNT, url, 200, {}, body) + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, body + b"\x00\x00\x02") + + def test_a_converse_stream_whose_stop_frame_names_no_reason_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, + eventstream_event("messageStart", {"role": "assistant"}) + + eventstream_event("messageStop", {}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}), + ) + + def test_an_invoke_stream_carrying_a_frame_that_is_not_a_chunk_is_not_recordable(self) -> None: + """Every frame of an invoke stream is a `chunk` holding one base64 event. + A frame that is not one carries an event this rule cannot read, so the + stream can no longer be judged complete.""" + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + eventstream_event("metadata", {"usage": {"totalTokens": 18}}) + + invoke_chunk({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_a_frame_claiming_no_length_is_rejected_rather_than_walked_forever(self) -> None: + """A frame length of zero never advances the cursor. Rejecting it is what + keeps a corrupt body from spinning the edge instead of answering.""" + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, b"\x00\x00\x00\x00" * 4) + + @pytest.mark.parametrize("url,terminator", [ + (INVOKE_STREAM_URL, invoke_chunk({"type": "message_stop"})), + (CONVERSE_STREAM_URL, eventstream_event("metadata", {"usage": {"totalTokens": 18}})), + ], ids=["invoke-stream", "converse-stream"]) + def test_a_delta_that_names_no_stop_reason_does_not_finish_a_stream( + self, url: str, terminator: bytes, + ) -> None: + """A `message_delta` arriving without its stop reason is the shape of a + turn the connection cut short partway through the delta itself.""" + head: Final = ( + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + invoke_chunk({"type": "message_delta", "delta": {}}) + ) + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, head + terminator) + + def test_an_invoke_chunk_that_is_not_base64_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + eventstream_event("chunk", {"bytes": "not base64 at all !!"}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_an_invoke_stream_missing_its_stop_reason_is_not_recordable(self) -> None: + assert not successful_response( + BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, + invoke_chunk({"type": "message_start", "message": {"id": "msg_bdrk_x"}}) + + invoke_chunk({"type": "message_stop"}), + ) + + def test_an_empty_stream_is_not_recordable(self) -> None: + for url in (CONVERSE_STREAM_URL, INVOKE_STREAM_URL): + assert not successful_response(BEDROCK_MOUNT, url, 200, {}, b"") + + def test_each_streaming_endpoint_is_held_to_its_own_grammar(self) -> None: + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 200, {}, INVOKE_STREAM_OK) + assert not successful_response(BEDROCK_MOUNT, INVOKE_STREAM_URL, 200, {}, CONVERSE_STREAM_OK) + + def test_a_stream_that_errored_before_it_started_is_not_recordable(self) -> None: + assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 503, {}, CONVERSE_STREAM_OK) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 698e78a5aa1..29d6e5ab4f2 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -4,7 +4,9 @@ The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored -Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, are not cacheable. They still cross the edge and are still re-signed, so they need the same IAM, but they always call the provider. AWS frames them as binary `vnd.amazon.eventstream` rather than SSE, and reading a terminal event out of that is what a completeness rule for them would need. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so most Bedrock traffic in the suite is not cached today +Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic + +Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way ## Request identity diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 0dee33f33c6..9ae89861f63 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -13,6 +13,7 @@ from types import MappingProxyType from typing import Final, Literal, Protocol from urllib.parse import urlsplit +from botocore.eventstream import EventStreamBuffer, ParserError from e2e_http import ( NetworkError, StreamChunk, @@ -36,6 +37,19 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" +BEDROCK_CONVERSE_SUFFIX: Final = "/converse" +BEDROCK_INVOKE_SUFFIX: Final = "/invoke" +BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream" +BEDROCK_INVOKE_STREAM_SUFFIX: Final = "/invoke-with-response-stream" +BEDROCK_SUFFIXES: Final = ( + BEDROCK_CONVERSE_SUFFIX, + BEDROCK_INVOKE_SUFFIX, + BEDROCK_CONVERSE_STREAM_SUFFIX, + BEDROCK_INVOKE_STREAM_SUFFIX, +) +EVENTSTREAM_PRELUDE_BYTES: Final = 4 +EVENT_TYPE_HEADER: Final = ":event-type" +EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -145,7 +159,7 @@ def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> return False path: Final = urlsplit(url).path if is_bedrock(mount): - return path.startswith("/model/") and path.endswith(("/converse", "/invoke")) + return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES) return path in OPENAI_JSON_PATHS @@ -176,16 +190,7 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": return events[-1] == "[DONE]" and "[DONE]" not in events[:-1] and complete_chat_stream(values) - return ( - "[DONE]" not in events - and isinstance(values[0], dict) and values[0].get("type") == "message_start" - and isinstance(values[-1], dict) and values[-1].get("type") == "message_stop" - and any( - isinstance(value, dict) and value.get("type") == "message_delta" - and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str) - for value in values - ) - ) + return "[DONE]" not in events and complete_anthropic_stream(values) try: value: Final = JSON_VALUE.validate_json(body) except ValidationError: @@ -214,14 +219,19 @@ def complete_bedrock_response(url: str, body: bytes) -> bool: """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an Anthropic model answers the Anthropic message shape. Either way a truncated or error body is missing the terminator field, which is what makes it safe to - record. The streaming variants never reach here: they are not cacheable.""" + record.""" + path: Final = urlsplit(url).path + if path.endswith(BEDROCK_CONVERSE_STREAM_SUFFIX): + return complete_converse_stream(body) + if path.endswith(BEDROCK_INVOKE_STREAM_SUFFIX): + return complete_invoke_stream(body) try: value: Final = JSON_VALUE.validate_json(body) except ValidationError: return False if not isinstance(value, dict) or "message" in value: return False - if urlsplit(url).path.endswith("/converse"): + if path.endswith(BEDROCK_CONVERSE_SUFFIX): return isinstance(value.get("output"), dict) and isinstance(value.get("stopReason"), str) return ( value.get("type") == "message" @@ -230,6 +240,111 @@ def complete_bedrock_response(url: str, body: bytes) -> bool: ) +def whole_eventstream_messages(body: bytes) -> bool: + """Whether the body is exactly a whole number of eventstream messages. + + A dropped connection is the failure this catches, and it has to be caught + here: botocore yields the messages it did receive and silently discards a + trailing partial one, so a stream cut a single byte short parses clean. Each + message declares its own total length in its first four bytes, so walking + those is enough to tell a complete body from a cut one.""" + offset = 0 # rebind-ok: a cursor walking the declared frame lengths + while offset + EVENTSTREAM_PRELUDE_BYTES <= len(body): + total: int = int.from_bytes(body[offset : offset + EVENTSTREAM_PRELUDE_BYTES], "big") + if total <= 0 or offset + total > len(body): + return False + offset += total + return offset == len(body) + + +def eventstream_events(body: bytes) -> tuple[tuple[str, JsonValue], ...] | None: + """The stream's (event type, decoded payload) pairs, or None if it is not a + complete, uncorrupted stream. + + botocore validates both CRCs and raises ``ParserError`` rather than decoding + corruption into something plausible. A failure that began after Bedrock had + already answered 200 arrives as an ``exception`` frame in place of the + terminator, so it is the terminator rules below that reject it and this does + not need to inspect ``:message-type`` as well.""" + if not body or not whole_eventstream_messages(body): + return None + buffer: Final = EventStreamBuffer() + buffer.add_data(body) + try: + return tuple( + (event_type(event.headers), JSON_VALUE.validate_json(event.payload)) + for event in buffer + ) + except (ParserError, ValidationError, ValueError): + return None + + +def event_type(headers: object) -> str: + """botocore's eventstream headers come back untyped, so the one header this + reads is validated into a string rather than trusted.""" + parsed: Final = EVENTSTREAM_HEADERS.validate_python(headers) + return parsed.get(EVENT_TYPE_HEADER, "") + + +def complete_converse_stream(body: bytes) -> bool: + """ConverseStream ends with ``metadata``, not with ``messageStop``. + + Requiring the metadata frame rather than the stop frame is deliberate: it + carries the token usage litellm prices the call from, so a stream cut between + the two still names a stop reason but would replay as a free call.""" + events: Final = eventstream_events(body) + if not events or events[-1][0] != "metadata": + return False + return any( + event_type == "messageStop" and isinstance(payload, dict) and isinstance(payload.get("stopReason"), str) + for event_type, payload in events + ) + + +def complete_invoke_stream(body: bytes) -> bool: + """InvokeModelWithResponseStream wraps the ordinary Anthropic event grammar + in ``chunk`` frames, one base64 payload each, so it is held to the same + terminator rule as the Anthropic SSE path. A frame Bedrock sends instead of a + chunk, an exception among them, carries no such payload and fails the rule + without the frame type needing to be read.""" + events: Final = eventstream_events(body) + if not events: + return False + values: Final = tuple(invoke_chunk_value(payload) for _, payload in events) + return all(value is not None for value in values) and complete_anthropic_stream(values) + + +def invoke_chunk_value(payload: JsonValue) -> JsonValue | None: + """The Anthropic event inside one ``chunk`` frame, or None for a frame that + carries no readable one.""" + if not isinstance(payload, dict) or not isinstance(encoded := payload.get("bytes"), str): + return None + try: + return JSON_VALUE.validate_json(base64.b64decode(encoded, validate=True)) + except (ValidationError, ValueError): + return None + + +def complete_anthropic_stream(values: tuple[JsonValue, ...]) -> bool: + """The Anthropic event grammar, shared by the SSE mounts and by Bedrock's + invoke stream, which carries the same events inside eventstream frames. A + ``message_delta`` naming a stop reason is what separates a finished turn from + one the connection cut short.""" + if not values: + return False + first: Final = values[0] + last: Final = values[-1] + return ( + isinstance(first, dict) and first.get("type") == "message_start" + and isinstance(last, dict) and last.get("type") == "message_stop" + and any( + isinstance(value, dict) and value.get("type") == "message_delta" + and isinstance(delta := value.get("delta"), dict) and isinstance(delta.get("stop_reason"), str) + for value in values + ) + ) + + def complete_responses_stream(values: tuple[JsonValue, ...]) -> bool: """The Responses API streams typed events and ends with ``response.completed``. A run that failed, was cancelled, or ran out of tokens ends with a different From 4d596082deecbf414abac3ee33e465717fa1e0e7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:20:32 +0000 Subject: [PATCH 084/187] feat(guardrails): release buffered stream chunks after each passing scan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/bedrock_guardrails.py | 3 + .../crowdstrike_aidr/__init__.py | 2 + .../crowdstrike_aidr/crowdstrike_aidr.py | 6 + .../unified_guardrail/unified_guardrail.py | 61 ++++---- .../guardrails/guardrail_initializers.py | 1 + litellm/types/guardrails.py | 6 + .../guardrail_hooks/crowdstrike_aidr.py | 8 + .../test_bedrock_guardrails.py | 3 + .../guardrail_hooks/test_crowdstrike_aidr.py | 13 ++ .../test_streaming_buffer_until_moderated.py | 139 +++++++++++++++++- 10 files changed, 215 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 2c407d91a48..db68c7b5c3f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -248,6 +248,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): streaming_buffer_until_moderated: bool | None = None, streaming_sampling_rate: int | None = None, streaming_end_of_stream_only: bool | None = None, + streaming_buffer_release_on_scan: bool | None = None, **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -258,6 +259,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): "streaming_buffer_until_moderated": streaming_buffer_until_moderated, "streaming_sampling_rate": streaming_sampling_rate, "streaming_end_of_stream_only": streaming_end_of_stream_only, + "streaming_buffer_release_on_scan": streaming_buffer_release_on_scan, } ) ) @@ -321,6 +323,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.streaming_buffer_until_moderated = streaming_params.streaming_buffer_until_moderated self.streaming_sampling_rate = streaming_params.streaming_sampling_rate self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only + self.streaming_buffer_release_on_scan = streaming_params.streaming_buffer_release_on_scan def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: super().update_in_memory_litellm_params(litellm_params) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py index c88e6e97a96..59f02817e5f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py @@ -23,6 +23,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" event_hook=litellm_params.mode, default_on=litellm_params.default_on, fail_on_error=litellm_params.fail_on_error, + streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated, + streaming_buffer_release_on_scan=streaming_params.streaming_buffer_release_on_scan, streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only, streaming_sampling_rate=streaming_params.streaming_sampling_rate, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 8fed1f906e5..924bbd2bc1a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -260,6 +260,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): api_key: str | None = None, api_base: str | None = None, fail_on_error: bool | None = True, + streaming_buffer_until_moderated: bool | None = None, + streaming_buffer_release_on_scan: bool | None = None, streaming_end_of_stream_only: bool | None = None, streaming_sampling_rate: int | None = None, async_handler: AsyncHTTPHandler | None = None, @@ -287,6 +289,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): CrowdStrikeAIDRGuardrailConfigModelOptionalParams( streaming_end_of_stream_only=streaming_end_of_stream_only, streaming_sampling_rate=streaming_sampling_rate, + streaming_buffer_until_moderated=streaming_buffer_until_moderated, + streaming_buffer_release_on_scan=streaming_buffer_release_on_scan, ) ) @@ -310,6 +314,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): ) def _set_streaming_params(self, streaming_params: CrowdStrikeAIDRGuardrailConfigModelOptionalParams) -> None: + self.streaming_buffer_until_moderated: bool = streaming_params.streaming_buffer_until_moderated or False + self.streaming_buffer_release_on_scan: bool = streaming_params.streaming_buffer_release_on_scan or False self.streaming_end_of_stream_only: bool = streaming_params.streaming_end_of_stream_only or False self.streaming_sampling_rate: int = streaming_params.streaming_sampling_rate or 5 diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index ee5cd7c4cb8..f231ac68c6a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -956,6 +956,7 @@ class UnifiedLLMGuardrails(CustomLogger): buffer_until_moderated: bool = _streaming_flag( "streaming_buffer_until_moderated", buffer_until_moderated_default ) + release_on_scan: bool = _streaming_flag("streaming_buffer_release_on_scan", False) if ( buffer_until_moderated @@ -972,7 +973,7 @@ class UnifiedLLMGuardrails(CustomLogger): # Buffering can only moderate the assembled response, so it always # defers to end-of-stream. - if buffer_until_moderated: + if buffer_until_moderated and not release_on_scan: end_of_stream_only = True if guardrail_to_apply is None: @@ -1026,6 +1027,7 @@ class UnifiedLLMGuardrails(CustomLogger): chunk_counter = 0 responses_so_far: Final[list[object]] = [] responses_yielded: Final[list[object]] = [] + withheld_items: Final[list[object]] = [] # mutable-ok: streaming window must be released incrementally pending_end_of_stream_items: Final[list[object]] = [] # Whether any real response chunk has been forwarded to the client. # Drives how a block terminates the stream: continue the in-progress @@ -1069,9 +1071,13 @@ class UnifiedLLMGuardrails(CustomLogger): chunks_yielded = True responses_yielded.append(item) yield item + else: + withheld_items.append(item) continue # Process chunk based on sampling rate + if buffer_until_moderated: + withheld_items.append(item) if chunk_counter % sampling_rate == 0: endpoint_translation = mappings[CallTypes(call_type)]() scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far) @@ -1081,9 +1087,16 @@ class UnifiedLLMGuardrails(CustomLogger): chunk_counter, guardrail_to_apply.guardrail_name, ) - chunks_yielded = True - responses_yielded.append(item) - yield item + if buffer_until_moderated: + for withheld_item in withheld_items: + chunks_yielded = True + responses_yielded.append(withheld_item) + yield withheld_item + withheld_items.clear() + else: + chunks_yielded = True + responses_yielded.append(item) + yield item continue verbose_proxy_logger.debug( @@ -1093,13 +1106,9 @@ class UnifiedLLMGuardrails(CustomLogger): guardrail_to_apply.guardrail_name, ) - # Deep-copy the current chunk before guardrail processing. - # process_output_streaming_response modifies responses_so_far - # in-place: it puts the combined guardrailed text in the first - # chunk and clears all subsequent chunks to "". Without this - # copy, yielding processed_items[-1] would yield an empty - # string, permanently losing this chunk's content. - original_item = copy.deepcopy(item) + original_items = ( + tuple(copy.deepcopy(withheld_items)) if buffer_until_moderated else (copy.deepcopy(item),) + ) try: await endpoint_translation.process_output_streaming_response( @@ -1144,13 +1153,16 @@ class UnifiedLLMGuardrails(CustomLogger): return if scan_key is not None: last_scan_key = scan_key - chunks_yielded = True - responses_yielded.append(original_item) - yield original_item + for original_item in original_items: + chunks_yielded = True + responses_yielded.append(original_item) + yield original_item + withheld_items.clear() else: - chunks_yielded = True - responses_yielded.append(item) - yield item + if not buffer_until_moderated: + chunks_yielded = True + responses_yielded.append(item) + yield item # Stream has ended - do final processing with all collected chunks if call_type is not None and CallTypes(call_type) in mappings: @@ -1162,14 +1174,13 @@ class UnifiedLLMGuardrails(CustomLogger): endpoint_translation = mappings[CallTypes(call_type)]() - # When buffering, snapshot the original chunks before moderation. - # A shallow copy suffices: end-of-stream - # process_output_streaming_response builds a separate assembled - # response (it does not mutate the individual chunks in place), and - # the chunks themselves are replayed verbatim -- so we only need to - # preserve the list, not clone every chunk (deepcopy would double - # peak memory for large responses). - buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None + buffered_items: Final = ( + tuple(copy.deepcopy(withheld_items)) + if buffer_until_moderated and release_on_scan and not end_of_stream_only + else tuple(withheld_items) + if buffer_until_moderated + else None + ) end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far) if _is_redundant_scan(end_scan_key, last_scan_key): verbose_proxy_logger.debug( diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 7858adeb55d..5dc9af2c081 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -44,6 +44,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated, streaming_sampling_rate=streaming_params.streaming_sampling_rate, streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only, + streaming_buffer_release_on_scan=streaming_params.streaming_buffer_release_on_scan, ) 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 92fe41ba717..6346e13f3ba 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -682,6 +682,12 @@ class BedrockGuardrailStreamingParams(BaseModel): "and the scan result lands in guardrail_information; a flagged response still ends the " "stream with a block message (disable_exception_on_block=true) or an error frame.", ) + streaming_buffer_release_on_scan: bool = Field( + default=False, + description="When buffering, scan the accumulated response every streaming_sampling_rate chunks " + "and release the withheld chunks once the scan passes, instead of holding everything to end of stream. " + "Flagged content is never released. Ignored when streaming_end_of_stream_only is true.", + ) @classmethod def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams": diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py index 6beca030a3a..df1caab6af6 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py @@ -4,6 +4,14 @@ from .base import GuardrailConfigModel class CrowdStrikeAIDRGuardrailConfigModelOptionalParams(BaseModel): + streaming_buffer_until_moderated: bool | None = Field( + default=None, + description="When True, withhold streamed chunks until moderation passes. Defaults to False when unset.", + ) + streaming_buffer_release_on_scan: bool | None = Field( + default=None, + description="When buffering, release withheld chunks after each passing scan. Defaults to False when unset.", + ) streaming_end_of_stream_only: bool | None = Field( default=None, description="If False (default when unset), post_call scans the accumulated streamed response every " 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 d173c5f5c70..dc275c912c8 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 @@ -5603,6 +5603,7 @@ def test_initialize_bedrock_wires_streaming_flags(): streaming_buffer_until_moderated=False, streaming_sampling_rate=3, streaming_end_of_stream_only=True, + streaming_buffer_release_on_scan=True, ), {"guardrail_name": "bedrock-streaming"}, ) @@ -5616,9 +5617,11 @@ def test_initialize_bedrock_wires_streaming_flags(): assert configured.streaming_buffer_until_moderated is False assert configured.streaming_sampling_rate == 3 assert configured.streaming_end_of_stream_only is True + assert configured.streaming_buffer_release_on_scan is True assert defaulted.streaming_buffer_until_moderated is True assert defaulted.streaming_sampling_rate == 5 assert defaulted.streaming_end_of_stream_only is False + assert defaulted.streaming_buffer_release_on_scan is False def test_initialize_bedrock_rejects_non_positive_sampling_rate(): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index 9849ad7ec88..a1aae119d56 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -1622,10 +1622,23 @@ def test_initialize_guardrail_rejects_unsupported_mode_instead_of_running_other_ def test_initialize_guardrail_defaults_streaming_params() -> None: handler = _initialize_from_config(mode="post_call") + assert handler.streaming_buffer_until_moderated is False + assert handler.streaming_buffer_release_on_scan is False assert handler.streaming_end_of_stream_only is False assert handler.streaming_sampling_rate == 5 +def test_initialize_guardrail_forwards_buffer_streaming_params() -> None: + handler = _initialize_from_config( + mode="post_call", + streaming_buffer_until_moderated=True, + streaming_buffer_release_on_scan=True, + ) + + assert handler.streaming_buffer_until_moderated is True + assert handler.streaming_buffer_release_on_scan is True + + @pytest.mark.parametrize( "configured", [ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py index 2b163ee5233..89d51866533 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py @@ -11,7 +11,7 @@ released unchanged after moderation passes. """ import json -from typing import Any, List, Literal, Optional +from typing import Any, AsyncGenerator, List, Literal, Optional import pytest @@ -23,7 +23,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import Delta, GenericGuardrailAPIInputs, ModelResponseStream, StreamingChoices BLOCK_MESSAGE = "Blocked by policy: this response was withheld." ORIGINAL_MARKER = "ORIGINAL-SECRET-ANSWER" @@ -60,6 +60,41 @@ class _PassingGuardrail(CustomGuardrail): return inputs +class _CountingPassingGuardrail(_PassingGuardrail): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.scan_count = 0 + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.scan_count += 1 + return inputs + + +class _SecondScanBlockingGuardrail(_CountingPassingGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.scan_count += 1 + if self.scan_count == 2: + raise ModifyResponseException( + message=BLOCK_MESSAGE, + model="gpt-4", + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + return inputs + + def _sse_event(event_type: str, data: dict) -> bytes: return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() @@ -115,6 +150,67 @@ def _decode(chunks: List[Any]) -> str: return "".join(c.decode() if isinstance(c, bytes) else str(c) for c in chunks) +def _chat_chunk(content: str = "", finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-windowed", + created=1724900000, + model="gpt-4", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content=content), + finish_reason=finish_reason, + ) + ], + ) + + +async def _windowed_chat_stream( + yielded_count: List[int], collected: List[Any], content_chunks: List[str] +) -> AsyncGenerator[ModelResponseStream, None]: + for content in content_chunks: + yielded_count.append(len(collected)) + yield _chat_chunk(content) + yielded_count.append(len(collected)) + yield _chat_chunk(finish_reason="stop") + + +async def _run_windowed( + guardrail: CustomGuardrail, + content_chunks: List[str], + end_of_stream_only: bool = False, +) -> tuple[List[Any], List[int]]: + guardrail.streaming_buffer_until_moderated = True + guardrail.streaming_buffer_release_on_scan = True + guardrail.streaming_end_of_stream_only = end_of_stream_only + guardrail.streaming_sampling_rate = 2 + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions") + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: List[Any] = [] + yielded_count: List[int] = [] + async for chunk in unified.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_windowed_chat_stream(yielded_count, collected, content_chunks), + request_data=request_data, + ): + collected.append(chunk) + return collected, yielded_count + + +def _chat_text(chunks: List[Any]) -> str: + return "".join( + choice.delta.content or "" + for chunk in chunks + if isinstance(chunk, ModelResponseStream) + for choice in chunk.choices + ) + + async def _run(guardrail: CustomGuardrail) -> str: # Rubrik's real config: end-of-stream-only moderation. Without buffering # this releases every chunk before moderation runs (content leaks on @@ -159,6 +255,45 @@ async def test_buffered_clean_releases_all_content(): assert BLOCK_MESSAGE not in raw +@pytest.mark.asyncio +async def test_windowed_buffer_releases_after_each_passing_scan(): + guardrail = _CountingPassingGuardrail(guardrail_name="windowed-pass", event_hook="post_call") + content_chunks = ["one ", "two ", "three ", "four ", "five ", "six "] + + collected, yielded_count = await _run_windowed(guardrail, content_chunks) + + assert yielded_count[2] >= 2 + assert yielded_count == [0, 0, 2, 2, 4, 4, 6] + assert _chat_text(collected) == "".join(content_chunks) + assert guardrail.scan_count > 1 + + +@pytest.mark.asyncio +async def test_windowed_buffer_drops_blocked_window(): + guardrail = _SecondScanBlockingGuardrail(guardrail_name="windowed-block", event_hook="post_call") + content_chunks = ["one ", "two ", "MARKER ", "four ", "five ", "six "] + + collected, _ = await _run_windowed(guardrail, content_chunks) + raw = _decode(collected) + + assert _chat_text(collected) == "one two " + assert "MARKER" not in raw + assert BLOCK_MESSAGE in raw + assert '"error"' not in raw + + +@pytest.mark.asyncio +async def test_windowed_buffer_with_explicit_end_of_stream_only_stays_fully_buffered(): + guardrail = _CountingPassingGuardrail(guardrail_name="windowed-eos", event_hook="post_call") + content_chunks = ["one ", "two ", "three ", "four ", "five ", "six "] + + collected, yielded_count = await _run_windowed(guardrail, content_chunks, end_of_stream_only=True) + + assert yielded_count == [0, 0, 0, 0, 0, 0, 0] + assert _chat_text(collected) == "".join(content_chunks) + assert guardrail.scan_count == 1 + + @pytest.mark.asyncio async def test_buffered_mode_disabled_for_content_rewriting_guardrail(): """Buffered replay yields the withheld *original* chunks verbatim, which From e64552867a1a3efd3f1ca88b94d94b5a982f8727 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:21:36 +0000 Subject: [PATCH 085/187] refactor(guardrails): drop stale buffering comment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/unified_guardrail/unified_guardrail.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index f231ac68c6a..70dc6eec24c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -971,8 +971,6 @@ class UnifiedLLMGuardrails(CustomLogger): ) buffer_until_moderated = False - # Buffering can only moderate the assembled response, so it always - # defers to end-of-stream. if buffer_until_moderated and not release_on_scan: end_of_stream_only = True From 8a553ceb58887c8aa2aa24c7cfb75ce662d2e321 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 07:38:44 -0700 Subject: [PATCH 086/187] feat(e2e): mount Gemini on the provider cache Gemini needs none of the machinery Bedrock needed. litellm composes {api_base}/models/{model}:{endpoint} from a custom api_base, so a plain path-prefixed mount reaches it, and the credential travels as a static x-goog-api-key header that no host rewrite invalidates. Nothing is re-signed and nothing leaves the cache key, so a recording still cannot cross credentials. A finished turn names a finishReason on every candidate and reports usageMetadata. The reason is read as a string rather than compared to STOP: MAX_TOKENS and the safety reasons end a turn just as finally, and rejecting them would send every one of them upstream forever. Streaming is the half worth care. Gemini repeats usageMetadata on every chunk and names a finishReason only on the last, so the terminator is the final event rather than any event, and a stream the connection cut short ends on a chunk carrying usage and no reason. The mount's upstream base carries the API version, so the path the rules see is /v1beta/models/..., not the one the proxy sent. The first version of this anchored the rule at the start of that path, which passed every test against a stub with no version prefix and would have cached nothing at all in a real run. Caught by replaying the rules over responses captured from live gemini-2.5-flash, which is also why the tests now mount their stub under the version prefix. Vertex stays unmounted and is a separate provider here: litellm grafts the default Vertex path onto an api_base only when that api_base has no path of its own, so Vertex needs a root-mounted edge on its own port. --- .../test_provider_cache.py | 154 +++++++++++++++++- tests/e2e/PROVIDER_CACHE.md | 16 +- tests/e2e/provider_cache.py | 46 ++++++ tests/e2e/provider_cache_routing.py | 3 +- tests/e2e/provider_edge.py | 1 + 5 files changed, 215 insertions(+), 5 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 4c131434ecd..520317a0678 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -861,7 +861,7 @@ def test_anthropic_stream_requires_start_finish_and_stop() -> None: assert not successful_response("anthropic", url, 200, headers, start + finish) -@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")]) +@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", ""), ("gemini", "")]) def test_normal_registration_routes_supported_providers(provider: str, suffix: str) -> None: params: Final = LiteLLMParamsBody(model=f"{provider}/test", api_key="os.environ/SYNTHETIC_KEY", timeout=12) routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) @@ -873,6 +873,8 @@ def test_normal_registration_routes_supported_providers(provider: str, suffix: s @pytest.mark.parametrize("params", [ LiteLLMParamsBody(model="bedrock/test"), LiteLLMParamsBody(model="azure/test"), + LiteLLMParamsBody(model="vertex_ai/gemini-2.5-flash"), + LiteLLMParamsBody(model="gemini/gemini-2.5-flash", api_base="https://custom.invalid"), LiteLLMParamsBody(model="openai/test", api_base="https://custom.invalid/v1"), LiteLLMParamsBody(model="openai/test", api_base=""), LiteLLMParamsBody(model="openai/test", litellm_credential_name="named-credential"), @@ -1241,3 +1243,153 @@ class TestBedrockStreams: def test_a_stream_that_errored_before_it_started_is_not_recordable(self) -> None: assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 503, {}, CONVERSE_STREAM_OK) + + +GEMINI_MODEL: Final = "gemini-2.5-flash" +GEMINI_API_VERSION: Final = "/v1beta" +GEMINI_GENERATE_PATH: Final = f"/models/{GEMINI_MODEL}:generateContent" +GEMINI_STREAM_PATH: Final = f"/models/{GEMINI_MODEL}:streamGenerateContent" +GEMINI_USAGE: Final = {"promptTokenCount": 7, "candidatesTokenCount": 1, "totalTokenCount": 25} + + +def gemini_body(finish_reason: str | None, usage: bool = True, candidates: bool = True) -> JsonValue: + candidate: Final[dict[str, JsonValue]] = {"content": {"parts": [{"text": "OK"}], "role": "model"}, "index": 0} + return { + "candidates": [{**candidate, "finishReason": finish_reason} if finish_reason else candidate] + if candidates else [], + **({"usageMetadata": GEMINI_USAGE} if usage else {}), + "modelVersion": GEMINI_MODEL, + } + + +def gemini_unary(finish_reason: str | None = "STOP", usage: bool = True, candidates: bool = True) -> bytes: + return json.dumps(gemini_body(finish_reason, usage, candidates)).encode() + + +def gemini_stream(*finish_reasons: str | None) -> bytes: + return b"".join( + b"data: " + json.dumps(gemini_body(reason)).encode() + b"\r\n\r\n" for reason in finish_reasons + ) + + +@contextmanager +def gemini_edge(cache: CacheEdge, provider: Provider, path: str) -> Generator[str, None, None]: + upstream: Final = f"http://127.0.0.1:{provider.server_port}{GEMINI_API_VERSION}" + running: Final = start_provider_edge(cache, mounts={"gemini": upstream}) + try: + yield running.edge.api_base("gemini") + path + finally: + running.shutdown() + + +class TestGemini: + """Gemini reaches the edge by path prefix alone: litellm composes + `{api_base}/models/{model}:{endpoint}` and sends a static `x-goog-api-key`, + so nothing has to be re-signed and nothing leaves the cache key. The response + grammar is its own though, and the streaming one is the interesting half: every + chunk repeats `usageMetadata`, so only `finishReason` on the last chunk + separates a finished turn from a dropped connection.""" + + @pytest.mark.parametrize("path,response", [ + (GEMINI_GENERATE_PATH, gemini_unary()), + (GEMINI_STREAM_PATH, gemini_stream(None, None, "STOP")), + ], ids=["generate", "stream"]) + def test_a_finished_turn_replays_on_the_next_run( + self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, + ) -> None: + provider.stream = path == GEMINI_STREAM_PATH + provider.response = response + for _ in range(2): + with gemini_edge(cache_edge(store), provider, path) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("reason", ["MAX_TOKENS", "SAFETY", "RECITATION"]) + def test_a_turn_the_provider_ended_for_its_own_reasons_is_still_finished( + self, store: RedisResponseStore, provider: Provider, reason: str, + ) -> None: + """Reading `finishReason` as a string rather than comparing it to STOP is + deliberate. A turn cut off by the token limit or a safety filter is over, + and rejecting those would send every one of them upstream forever.""" + provider.response = gemini_unary(reason) + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: + assert call(url, MARKED).body == provider.response + assert len(provider.hits) == 1 + + @pytest.mark.parametrize("response", [ + gemini_unary(None), + gemini_unary("STOP", usage=False), + gemini_unary("STOP", candidates=False), + b'{"error":{"code":400,"message":"API key not valid","status":"INVALID_ARGUMENT"}}', + ], ids=["no-finish-reason", "no-usage", "no-candidates", "error-body"]) + def test_an_unfinished_or_failed_turn_never_enters_the_cache( + self, store: RedisResponseStore, provider: Provider, response: bytes, + ) -> None: + provider.response = response + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("response", [ + gemini_stream(None, None), + gemini_stream("STOP", None), + gemini_stream(), + ], ids=["cut-before-the-reason", "reason-then-another-chunk", "empty"]) + def test_a_stream_that_never_named_a_reason_calls_the_provider_every_time( + self, store: RedisResponseStore, provider: Provider, response: bytes, + ) -> None: + provider.stream = True + provider.response = response + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_STREAM_PATH) as url: + assert call(url, MARKED).body == response + assert len(provider.hits) == 2 + + def test_a_response_whose_candidates_did_not_all_finish_is_not_recordable( + self, store: RedisResponseStore, provider: Provider, + ) -> None: + """A request for more than one candidate is answered by more than one, and + the turn is over only when every one of them names a reason. Holding the + whole list to that rule rather than its first entry is what keeps a + half-finished answer from being stored and replayed as a finished one.""" + finished: Final = json.loads(gemini_unary("STOP"))["candidates"][0] + unfinished: Final = json.loads(gemini_unary(None))["candidates"][0] + provider.response = json.dumps( + {"candidates": [finished, {**unfinished, "index": 1}], "usageMetadata": GEMINI_USAGE} + ).encode() + for _ in range(2): + with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: + assert call(url, MARKED).body == provider.response + assert len(provider.hits) == 2 + + @pytest.mark.parametrize("path,cacheable", [ + (GEMINI_GENERATE_PATH, True), + (GEMINI_STREAM_PATH, True), + (f"/models/{GEMINI_MODEL}:countTokens", False), + (f"/models/{GEMINI_MODEL}:embedContent", False), + ("/v1/chat/completions", False), + (f"/files/{GEMINI_MODEL}:generateContent", False), + ]) + @pytest.mark.parametrize("version", ["", GEMINI_API_VERSION], ids=["bare", "versioned"]) + def test_only_the_generate_endpoints_are_cacheable(self, version: str, path: str, cacheable: bool) -> None: + """The mount's upstream base carries the API version, so the path the cache + sees is the upstream one and starts `/v1beta`. A rule anchored at the start + of the path would pass every test against a stub with no version prefix and + then cache nothing at all in a real run.""" + assert cacheable_endpoint("gemini", "POST", f"https://gemini.invalid{version}{path}", MARKED) is cacheable + + def test_the_bodies_these_tests_build_match_a_real_gemini_response(self) -> None: + """The shapes above are hand-built so a test can express the turn it means. + This holds them to the fields a live `generativelanguage.googleapis.com` + answer carries, captured 2026-09-16 against gemini-2.5-flash.""" + captured: Final = json.loads( + '{"candidates":[{"content":{"parts":[{"text":"OK"}],"role":"model"},"finishReason":"STOP",' + '"index":0}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":1,' + '"totalTokenCount":25},"modelVersion":"gemini-2.5-flash","responseId":"1J6qauKFI8ut1MkPgNjI4AI"}' + ) + built: Final = json.loads(gemini_unary()) + assert captured.keys() >= built.keys() + assert captured["candidates"][0].keys() >= built["candidates"][0].keys() + assert successful_response("gemini", GEMINI_GENERATE_PATH, 200, {}, json.dumps(captured).encode()) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 29d6e5ab4f2..e2f8030074d 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -1,13 +1,23 @@ # Shared provider-response cache -`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live +`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI, Anthropic and Gemini model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live -The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, and for `models/{model}:generateContent` and `:streamGenerateContent` on the Gemini mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way +## Gemini + +Gemini needs nothing that Bedrock needed. litellm composes `{api_base}/models/{model}:{endpoint}` from a custom api_base, so a path-prefixed mount reaches it, and the credential travels as a static `x-goog-api-key` header that no host rewrite invalidates. Nothing is re-signed and nothing is excluded from the key, so a recording still cannot cross credentials + +The mount's upstream base carries the API version, which is the one detail worth remembering: the path the cache rules see is the upstream one, `/v1beta/models/...`, not the one the proxy sent. A rule anchored at the start of that path would look right against a local stub and then cache nothing at all in a real run + +A finished turn names a `finishReason` on every candidate and reports `usageMetadata`. The reason is read as a string rather than compared to `STOP`, because `MAX_TOKENS` and the safety reasons end a turn just as finally and rejecting them would send every one of them upstream forever. Streaming is the more interesting half: Gemini repeats `usageMetadata` on every chunk and names a `finishReason` only on the last one, so the terminator is the final event rather than any event, and a stream the connection cut short ends on a chunk with usage and no reason + +Vertex is not mounted. litellm grafts the default Vertex path onto an api_base only when that api_base has no path of its own, so a Vertex mount needs a root-mounted edge on its own port rather than a path prefix. Gemini and Vertex are separate providers in litellm and the Gemini mount does not cover Vertex deployments + ## Request identity A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is @@ -30,7 +40,7 @@ Only deployments that carry no AWS identity of their own route to the edge. A de Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here -Vertex and Gemini are not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm +Vertex is not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm. Gemini is a separate provider there and does have a working path-prefixed form, so it is mounted; see the Gemini section Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 9ae89861f63..7bba93321b5 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -37,6 +37,10 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" +GEMINI_MOUNT: Final = "gemini" +GEMINI_MODELS_SEGMENT: Final = "/models" +GEMINI_GENERATE_SUFFIX: Final = ":generateContent" +GEMINI_STREAM_SUFFIX: Final = ":streamGenerateContent" BEDROCK_CONVERSE_SUFFIX: Final = "/converse" BEDROCK_INVOKE_SUFFIX: Final = "/invoke" BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream" @@ -154,12 +158,21 @@ def is_bedrock(mount: str) -> bool: return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX +def is_gemini(mount: str) -> bool: + return mount == GEMINI_MOUNT + + def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool: if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES: return False path: Final = urlsplit(url).path if is_bedrock(mount): return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES) + if is_gemini(mount): + collection, _, resource = path.rpartition("/") + return collection.endswith(GEMINI_MODELS_SEGMENT) and resource.endswith( + (GEMINI_GENERATE_SUFFIX, GEMINI_STREAM_SUFFIX) + ) return path in OPENAI_JSON_PATHS @@ -186,6 +199,8 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, for value in values ): return False + if is_gemini(mount): + return complete_gemini_stream(values) if urlsplit(url).path == "/v1/responses": return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": @@ -197,6 +212,8 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not isinstance(value, dict) or value.get("error") is not None: return False + if is_gemini(mount): + return complete_gemini_candidates(value) path: Final = urlsplit(url).path if path == "/v1/messages": return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str) @@ -215,6 +232,35 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, ) +def complete_gemini_candidates(value: Mapping[str, JsonValue]) -> bool: + """A finished Gemini turn names a ``finishReason`` on every candidate and + reports the usage litellm prices the call from. ``finishReason`` is read as a + string rather than compared to ``STOP`` because ``MAX_TOKENS`` and the safety + reasons end a turn just as finally, and a cache that rejected them would send + every one of them upstream forever.""" + candidates: Final = value.get("candidates") + return ( + isinstance(value.get("usageMetadata"), dict) + and isinstance(candidates, list) + and bool(candidates) + and all( + isinstance(candidate, dict) and isinstance(candidate.get("finishReason"), str) + for candidate in candidates + ) + ) + + +def complete_gemini_stream(values: tuple[JsonValue, ...]) -> bool: + """Gemini repeats ``usageMetadata`` on every chunk but names a + ``finishReason`` only on the last one, so the terminator is the final event + rather than any event. A stream the connection cut short ends on a chunk that + carries usage and no reason, which is exactly what this rejects.""" + if not values: + return False + last: Final = values[-1] + return isinstance(last, dict) and complete_gemini_candidates(last) + + def complete_bedrock_response(url: str, body: bytes) -> bool: """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an Anthropic model answers the Anthropic message shape. Either way a truncated diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index f9775a2b152..c4b02beac2d 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -19,6 +19,7 @@ BEDROCK_EDGE_MODELS: Final = frozenset( } ) ENV_REFERENCE_PREFIX: Final = "os.environ/" +EDGE_PROVIDERS: Final = frozenset({"openai", "anthropic", "gemini"}) def bedrock_region(declared: str | None) -> str: @@ -81,7 +82,7 @@ def route_cache_model( return route_bedrock(params, base_for, mode) if mode == "realtime" or params.api_base is not None: return params - if provider not in {"openai", "anthropic"}: + if provider not in EDGE_PROVIDERS: return params base: Final = base_for(provider) if base is None: diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 2606b26fe99..219df55233a 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -104,6 +104,7 @@ EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { "openai": "https://api.openai.com", "anthropic": "https://api.anthropic.com", + "gemini": "https://generativelanguage.googleapis.com/v1beta", **{ f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com" for region in BEDROCK_REGIONS From 7006da9cde950981379214f3d2dd0f645d68995e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 07:47:31 -0700 Subject: [PATCH 087/187] feat(e2e): say why a response was not recorded Build 226 routed Bedrock streaming for the first time and rejected 62 of 220 misses on that mount, and the counters could not say why. A flat rejected count covers three unrelated things with opposite fixes: the consumer walking away mid-capture, a body that arrived whole and failed its endpoint's rule, and a provider that could not be reached. Each now also counts its own reason. A consumer that walks away was counting nothing at all. Abandoning the capture generator raises GeneratorExit at its yield, so neither branch of the old accounting ran and the miss simply vanished from the report, which is also why misses could exceed writes plus rejected with nothing to explain the gap. The decision moves into settle() so the generator's finally owns the accounting and an abandoned capture is counted like any other rejection. --- .../test_provider_cache.py | 39 ++++++++++++++ tests/e2e/PROVIDER_CACHE.md | 2 +- tests/e2e/provider_cache.py | 53 +++++++++++++------ 3 files changed, 78 insertions(+), 16 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 520317a0678..1271b20438a 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -516,6 +516,45 @@ def test_counters_attribute_every_outcome_to_its_mount( assert counts["mount:anthropic:rejected"] == 1 and "mount:openai:rejected" not in counts +def test_a_rejection_says_whether_the_body_was_cut_short_or_simply_unfinished( + store: RedisResponseStore, provider: Provider, +) -> None: + """One `rejected` count cannot tell a connection that dropped from a body the + provider finished sending and the rules turned down, and those have opposite + fixes: the first is the client going away mid-capture, the second is a grammar + the cache does not accept. A mount whose rejections are mostly one or the other + is a different problem, so the report has to be able to say which.""" + upstream: Final = f"http://127.0.0.1:{provider.server_port}" + cut_short: Final = cache_edge(store) + provider.stream = True + provider.truncated = True + provider.response = b'data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":"stop"}]}\n\ndata: [DONE]\n\n' + running: Final = start_provider_edge(cut_short, mounts={"openai": upstream}) + try: + forward("POST", running.edge.api_base("openai") + "/v1/chat/completions", + headers=HEADERS, body=MARKED, timeout=5) + finally: + running.shutdown() + + unfinished: Final = cache_edge(store) + provider.stream = False + provider.truncated = False + provider.response = b'{"choices":[{"index":0,"message":{"content":"hi"}}]}' + second: Final = start_provider_edge(unfinished, mounts={"openai": upstream}) + try: + call(second.edge.api_base("openai") + "/v1/chat/completions", MARKED) + finally: + second.shutdown() + + cut: Final = dict(cut_short.counters.counts) + turned_down: Final = dict(unfinished.counters.counts) + assert cut["mount:openai:rejected"] == 1 and turned_down["mount:openai:rejected"] == 1 + assert cut["mount:openai:rejected_cut_short"] == 1 + assert "mount:openai:rejected_incomplete" not in cut + assert turned_down["mount:openai:rejected_incomplete"] == 1 + assert "mount:openai:rejected_cut_short" not in turned_down + + EMBEDDING_SUCCESS: Final = ( b'{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2]}],' b'"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index e2f8030074d..894d9be4efa 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -54,7 +54,7 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_incomplete` (the body arrived whole and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 7bba93321b5..f528d08b720 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -52,6 +52,9 @@ BEDROCK_SUFFIXES: Final = ( BEDROCK_INVOKE_STREAM_SUFFIX, ) EVENTSTREAM_PRELUDE_BYTES: Final = 4 +CUT_SHORT: Final = "cut_short" +INCOMPLETE: Final = "incomplete" +UNREACHABLE: Final = "unreachable" EVENT_TYPE_HEADER: Final = ":event-type" EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) @@ -551,7 +554,7 @@ class CacheEdge: ) prepared: Final = prepare_forward(method, url, self.outbound(mount, method, url, headers, body), body) if isinstance(prepared, NetworkError): - self.count(mount, "rejected") + self.reject(mount, UNREACHABLE) return prepared identity: Final = request_identity( self.secret, test_key, method, url, self.keyed(mount, prepared.headers), body, @@ -575,7 +578,7 @@ class CacheEdge: return head if isinstance(head, NetworkError): self.store.release(key, capture_slot) - self.count(mount, "rejected") + self.reject(mount, UNREACHABLE) return head return StreamHead( head.status_code, head.headers, primed_steps(self.capture(mount, key, capture_slot, url, head)), @@ -585,25 +588,45 @@ class CacheEdge: self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, ) -> Generator[StreamStep, None, None]: capture: Final = ResponseCapture() + reason = CUT_SHORT # rebind-ok: a consumer that walks away never reaches the settle call below try: with closing(head.steps): yield StreamChunk(b"") for step in head.steps: yield step capture.observe(step) - chunks: Final = capture.chunks() if capture.eligible else () - headers: Final = { - name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS - } - if not capture.eligible or not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): - self.count(mount, "rejected") - return - response: Final = CachedResponse( - request_key=key, status_code=head.status_code, headers=headers, - chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks), - ) - published: Final = self.store.publish(key, lease, encode_response(self.secret, response)) - self.count(mount, "writes" if published else "write_failures") + reason = self.settle(mount, key, lease, url, head, capture) finally: + self.reject(mount, reason) self.store.release(key, lease) capture.buffer.close() + + def settle( + self, mount: str, key: str, lease: CaptureLease, url: str, head: StreamHead, capture: ResponseCapture, + ) -> str | None: + """None once the response is stored, otherwise the reason it was not.""" + if not capture.eligible: + return CUT_SHORT + headers: Final = { + name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS + } + chunks: Final = capture.chunks() + if not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): + return INCOMPLETE + response: Final = CachedResponse( + request_key=key, status_code=head.status_code, headers=headers, + chunks=tuple(base64.b64encode(chunk).decode("ascii") for chunk in chunks), + ) + published: Final = self.store.publish(key, lease, encode_response(self.secret, response)) + self.count(mount, "writes" if published else "write_failures") + return None + + def reject(self, mount: str, reason: str | None) -> None: + """A flat rejection count cannot separate a connection that went away from + a body the provider finished sending and the rules turned down, and the two + have opposite fixes. A mount whose rejections are nearly all one or the + other is a different problem, so the report has to be able to say which.""" + if reason is None: + return + self.count(mount, "rejected") + self.count(mount, f"rejected_{reason}") From fdb8e3533be56df5de394bdb155870008c50bf12 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 07:58:48 -0700 Subject: [PATCH 088/187] fix(mcp): validate credentials in existing request paths --- .../mcp_server/mcp_server_manager.py | 42 ++++----- .../mcp_server/openapi_to_mcp_generator.py | 40 +++++++-- .../outbound_credentials/adapter.py | 67 ++++++++++++++- .../proxy/_experimental/mcp_server/server.py | 1 - .../_experimental/mcp_server/upstream.py | 85 ------------------- .../proxy/_experimental/mcp_server/utils.py | 16 ---- .../mcp_server/test_mcp_hook_extra_headers.py | 1 - .../mcp_server/test_mcp_server_manager.py | 76 ++++++++++++++--- .../test_openapi_to_mcp_generator.py | 84 ++++++++++++++++++ 9 files changed, 265 insertions(+), 147 deletions(-) delete mode 100644 litellm/proxy/_experimental/mcp_server/upstream.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0254f79cbcc..6881956595c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -102,6 +102,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( UpstreamCredentialProvider, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + prepare_mcp_client, raise_public, raise_token_exchange_challenge, raise_user_oauth_challenge, @@ -132,7 +133,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) -from litellm.proxy._experimental.mcp_server.upstream import prepare_mcp_client, validate_openapi_credentials from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, MCPMissingUserEnvVarsError, @@ -2805,6 +2805,8 @@ class MCPServerManager: headers=headers, server_label=server.name or server.server_name or server.alias or server.server_id, relays_upstream_auth=server.is_client_forwarded_token, + auth_type=server.auth_type, + upstream_token_header=server.upstream_token_header, ) tool_func.__name__ = prefixed_tool_name tool_func.__doc__ = description @@ -4230,19 +4232,16 @@ class MCPServerManager: ) record_auth_resolution(server.server_id, AuthResolution.not_applicable) - return await prepare_mcp_client( - resolved_server, - MCPClient( - server_url="", # Not used for stdio - transport_type=transport, - auth_type=resolved_server.auth_type, - auth_value=auth_value, - timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), - stdio_config=stdio_config, - extra_headers=extra_headers, - sampling_callback=sampling_cb, - elicitation_callback=elicitation_cb, - ), + return MCPClient( + server_url="", # Not used for stdio + transport_type=transport, + auth_type=resolved_server.auth_type, + auth_value=auth_value, + timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), + stdio_config=stdio_config, + extra_headers=extra_headers, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, ) else: # For HTTP/SSE transports @@ -6200,7 +6199,6 @@ class MCPServerManager: mcp_auth_header: str | dict[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None, forwarded_headers: dict[str, str] | None, - caller_authorization: str | None = None, ) -> tuple[dict[str, str] | None, dict[str, str] | None]: """Resolve the gateway-owned upstream credential for a spec_path (OpenAPI) tool call. @@ -6224,12 +6222,9 @@ class MCPServerManager: """ spec: Final = to_server_spec(mcp_server) if spec is None: - stored_headers = ( - None - if oauth2_headers - else await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) - ) - validate_openapi_credentials(mcp_server, stored_headers, forwarded_headers, caller_authorization) + if oauth2_headers: + return None, forwarded_headers + stored_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) return stored_headers, forwarded_headers subject_token: str | None = None @@ -6248,9 +6243,7 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, extra_headers=forwarded_headers, ) - resolved_headers: Final = await _materialize_auth_headers(resolved_auth) - validate_openapi_credentials(mcp_server, resolved_headers, forwarded_headers, caller_authorization) - return resolved_headers, forwarded_headers + return await _materialize_auth_headers(resolved_auth), forwarded_headers async def _gather_openapi_tool_tasks( self, @@ -6376,7 +6369,6 @@ class MCPServerManager: mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, - caller_authorization=auth_header_value, ) async def _call_openapi_via_handler(): diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 66712e97a34..477d86ab436 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -20,7 +20,6 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( MCPOpenApiUpstreamError, MCPUpstreamAuthError, ) -from litellm.proxy._experimental.mcp_server.utils import merge_openapi_headers # Tool names emitted from OpenAPI specs must work across all major LLM providers. # OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to @@ -55,7 +54,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) -from litellm.types.mcp import credential_redirect_hook, custom_credential_slot +from litellm.types.mcp import MCPAuthType, credential_redirect_hook, custom_credential_slot class _OpenAPIJSONSchema(TypedDict, total=False): @@ -416,9 +415,26 @@ def _merge_openapi_tool_request_headers( Header names are compared case-insensitively so different casing cannot bypass the precedence rules. """ - return merge_openapi_headers( - static_headers, _request_extra_headers.get(), _request_auth_header.get(), _request_resolved_auth_headers.get() - ) + request_extra: Final = _request_extra_headers.get() or {} + static: Final = static_headers or {} + + static_lower_names: Final = {k.lower() for k in static} + effective_headers: dict[str, str] = {k: v for k, v in request_extra.items() if k.lower() not in static_lower_names} + effective_headers.update(static) + + override_auth: Final = _request_auth_header.get() + if override_auth: + for existing in [k for k in effective_headers if k.lower() == "authorization"]: + del effective_headers[existing] + effective_headers["Authorization"] = override_auth + + resolved_auth_headers: Final = _request_resolved_auth_headers.get() or {} + for name, value in resolved_auth_headers.items(): + for existing in [k for k in effective_headers if k.lower() == name.lower()]: + del effective_headers[existing] + effective_headers[name] = value + + return effective_headers def _raise_for_upstream_failure( @@ -455,6 +471,8 @@ def create_tool_function( headers: dict[str, str] | None = None, server_label: str | None = None, relays_upstream_auth: bool = False, + auth_type: MCPAuthType = None, + upstream_token_header: str | None = None, ): """Create a tool function for an OpenAPI operation. @@ -487,6 +505,18 @@ def create_tool_function( by using **kwargs instead of named parameters. """ effective_headers: Final = _merge_openapi_tool_request_headers(headers) + if auth_type is not None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_public, + validate_static_credential, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok + + match validate_static_credential(auth_type, effective_headers, upstream_token_header): + case Error(error): + raise_public(error) + case Ok(): + pass # Build URL from base_url and path url = base_url + path diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index d25946d81d0..5358878a248 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -13,15 +13,17 @@ from __future__ import annotations import base64 import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Literal, NoReturn from fastapi import HTTPException from pydantic import SecretStr from typing_extensions import assert_never -from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials +from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( DEFAULT_CREDENTIAL_HEADER, ApiKeyConfig, @@ -39,7 +41,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( Subject, TokenExchangeConfig, ) -from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth +from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth, MCPAuthType, MCPTransport if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth @@ -385,3 +387,64 @@ def raise_token_exchange_challenge( detail="Unauthorized", headers={"WWW-Authenticate": www_authenticate}, ) + + +_STATIC_MODES: Final = frozenset( + (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization) +) + + +def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool: + if not value: + return False + if auth_type == MCPAuth.authorization or (auth_type == MCPAuth.api_key and name != "authorization"): + return True + if value.lower() in ("bearer", "basic", "token", "apikey"): + return False + if auth_type in (MCPAuth.bearer_token, MCPAuth.token): + scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token" + credential: Final = strip_auth_scheme(value, scheme).strip() + return bool(credential) and credential.lower() != scheme.lower() + if auth_type == MCPAuth.basic: + parts: Final = value.split(None, 1) + if len(parts) != 2 or parts[0].lower() != "basic": + return False + try: + decoded: Final = base64.b64decode(parts[1], validate=True).strip() + return b":" in decoded + except ValueError: + return False + return True + + +def validate_static_credential( + auth_type: MCPAuthType, + headers: Mapping[str, str], + upstream_token_header: str | None = None, +) -> Result[None, CredError]: + if auth_type not in _STATIC_MODES: + return Ok(None) + default_slot: Final = "X-API-Key" if auth_type == MCPAuth.api_key else "Authorization" + slots: Final = frozenset( + name.lower() + for name in ( + upstream_token_header or default_slot, + default_slot, + "Authorization", + ) + ) + values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) + if any(_usable_credential_value(auth_type, name, value) for name, value in values): + return Ok(None) + return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential")) + + +async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient: + if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio: + return client + request: Final = await client.prepare_request_auth() + match validate_static_credential(server.auth_type, request.headers, server.upstream_token_header): + case Error(error): + raise_public(error) + case Ok(): + return client diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a3aaada41f7..7feb1fd468d 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3141,7 +3141,6 @@ if MCP_AVAILABLE: mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, - caller_authorization=auth_header_value, ) _auth_token: Final = _request_auth_header.set(auth_header_value) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py deleted file mode 100644 index 66840db21ce..00000000000 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ /dev/null @@ -1,85 +0,0 @@ -from __future__ import annotations - -import base64 -from collections.abc import Mapping -from typing import Final - -from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme -from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import raise_public -from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result -from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError -from litellm.proxy._experimental.mcp_server.utils import merge_openapi_headers -from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransport -from litellm.types.mcp_server.mcp_server_manager import MCPServer - -_STATIC_MODES: Final = frozenset( - (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization) -) - - -def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool: - if not value: - return False - if auth_type == MCPAuth.authorization or (auth_type == MCPAuth.api_key and name != "authorization"): - return True - if value.lower() in ("bearer", "basic", "token", "apikey"): - return False - if auth_type in (MCPAuth.bearer_token, MCPAuth.token): - scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token" - credential: Final = strip_auth_scheme(value, scheme).strip() - return bool(credential) and credential.lower() != scheme.lower() - if auth_type == MCPAuth.basic: - parts: Final = value.split(None, 1) - if len(parts) != 2 or parts[0].lower() != "basic": - return False - try: - decoded: Final = base64.b64decode(parts[1], validate=True).strip() - return b":" in decoded - except ValueError: - return False - return True - - -def validate_static_credential(server: MCPServer, headers: Mapping[str, str]) -> Result[None, CredError]: - if server.auth_type not in _STATIC_MODES or server.transport == MCPTransport.stdio: - return Ok(None) - default_slot: Final = "X-API-Key" if server.auth_type == MCPAuth.api_key else "Authorization" - slots: Final = frozenset( - name.lower() - for name in ( - server.upstream_token_header or default_slot, - default_slot, - "Authorization", - ) - ) - values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) - if any(_usable_credential_value(server.auth_type, name, value) for name, value in values): - return Ok(None) - return Error(CredError.of_misconfigured(f"{server.auth_type} requires a usable upstream credential")) - - -async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient: - if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio: - return client - request: Final = await client.prepare_request_auth() - match validate_static_credential(server, request.headers): - case Error(error): - raise_public(error) - case Ok(): - return client - - -def validate_openapi_credentials( - server: MCPServer, - resolved_headers: Mapping[str, str] | None, - forwarded_headers: Mapping[str, str] | None, - caller_authorization: str | None, -) -> None: - headers: Final = merge_openapi_headers( - server.static_headers or {}, forwarded_headers, caller_authorization, resolved_headers - ) - match validate_static_credential(server, headers): - case Error(error): - raise_public(error) - case Ok(): - return diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index bea74d36b34..fb3eb06fd15 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -756,22 +756,6 @@ def build_env_var_setup_url(server_id: str) -> str: return f"{base}{path}" if base else path -def merge_openapi_headers( - static_headers: Mapping[str, str], - extra_headers: Mapping[str, str] | None, - caller_authorization: str | None, - resolved_headers: Mapping[str, str] | None, -) -> dict[str, str]: - sources: Final = ( - extra_headers or {}, - static_headers, - {"Authorization": caller_authorization} if caller_authorization else {}, - resolved_headers or {}, - ) - entries: Final = {name.lower(): (name, value) for source in sources for name, value in source.items()} - return dict(entries.values()) - - def merge_mcp_headers( *, extra_headers: Mapping[str, str] | None = None, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 5e3a26fb4ac..28faf375ab8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1391,7 +1391,6 @@ class TestOpenApiResolvedUpstreamAuth: mcp_auth_header="user-byok-key", user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="sk-user"), forwarded_headers=None, - caller_authorization="ApiKey user-byok-key", ) assert resolved is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 54add273c24..4c4c45162ca 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5,11 +5,13 @@ import logging import os import sys from datetime import datetime +from pathlib import Path from typing import Any, Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from respx import MockRouter from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, @@ -5127,7 +5129,8 @@ class TestMCPServerManager: captured: dict = {} def fake_create_tool_function( - path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False + path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False, + auth_type=None, upstream_token_header=None, ): captured["headers"] = headers captured["server_label"] = server_label @@ -5212,7 +5215,8 @@ class TestMCPServerManager: captured: dict = {} def fake_create_tool_function( - path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False + path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False, + auth_type=None, upstream_token_header=None, ): captured["headers"] = headers @@ -13471,6 +13475,41 @@ async def test_discovery_cache_returns_oversized_results_without_retaining_them( class TestProtectedCredentialPreparation: + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,credential", [ + (MCPAuth.bearer_token, None), + (MCPAuth.bearer_token, "Bearer"), + (MCPAuth.api_key, None), + (MCPAuth.basic, "Basic"), + ]) + @pytest.mark.parametrize("dispatch", ["managed", "local"]) + async def test_openapi_dispatch_rejects_unusable_effective_credentials( + self, tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthType, credential: str | None, dispatch: str, + ) -> None: + from litellm.proxy._experimental.mcp_server.server import _handle_local_mcp_tool + from litellm.proxy._experimental.mcp_server.utils import add_server_prefix_to_name, get_server_prefix + + spec_path: Final = tmp_path / "openapi.json" + spec_path.write_text(json.dumps({"openapi": "3.0.0", "info": {"title": "Auth", "version": "1"}, + "paths": {"/echo": {"get": {"operationId": "echo"}}}})) + server: Final = MCPServer( + server_id="dispatch-auth", name="dispatch-auth", url="https://upstream.example", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=credential, + ) + manager: Final = MCPServerManager() + await manager._register_openapi_tools(str(spec_path), server, server.url) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="unexpected success") + result: Final = ( + await manager._call_openapi_tool_handler(server, "echo", {}) + if dispatch == "managed" + else await _handle_local_mcp_tool(add_server_prefix_to_name("echo", get_server_prefix(server)), {}) + ) + assert result.isError is True + assert "requires a usable upstream credential" in result.content[0].text + assert destination.call_count == 0 + @pytest.mark.asyncio @pytest.mark.parametrize("transport", [MCPTransport.http, MCPTransport.sse]) @pytest.mark.parametrize("client_secret", [None, ""]) @@ -13523,7 +13562,7 @@ class TestProtectedCredentialPreparation: assert client._get_auth_headers() == headers @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.api_key, MCPAuth.bearer_token]) + @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange]) async def test_openapi_protected_auth_rejects_missing_credentials(self, auth_type: MCPAuthType) -> None: server = MCPServer( server_id="openapi-empty", name="openapi-empty", url="https://upstream.example/mcp", @@ -13594,22 +13633,35 @@ class TestProtectedCredentialPreparation: ({"X-API-Key": "static"}, {"Authorization": ""}, None), ]) async def test_openapi_static_credentials_remain_supported( - self, static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None + self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None ) -> None: - server = MCPServer(server_id="openapi-static", name="openapi-static", url="https://upstream.example", - transport=MCPTransport.http, auth_type=MCPAuth.api_key, static_headers=static) - resolved, retained = await MCPServerManager().resolve_openapi_upstream_auth( - mcp_server=server, oauth2_headers=None, raw_headers=None, mcp_auth_header=None, - user_api_key_auth=None, forwarded_headers=forwarded, caller_authorization=caller, + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_auth_header, _request_extra_headers, create_tool_function, ) - assert resolved is None - assert retained == forwarded + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.api_key, + ) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + caller_token: Final = _request_auth_header.set(caller) + extra_token: Final = _request_extra_headers.set(forwarded) + try: + assert await tool() == "authenticated" + sent: Final = destination.calls.last.request.headers + assert sent.get("x-api-key") == static.get("X-API-Key", (forwarded or {}).get("X-API-Key")) + if caller: + assert sent["authorization"] == caller + assert destination.call_count == 1 + finally: + _request_auth_header.reset(caller_token) + _request_extra_headers.reset(extra_token) @pytest.mark.asyncio async def test_static_resolution_cancellation_closes_flow(self) -> None: from collections.abc import AsyncGenerator from litellm.experimental_mcp_client.client import MCPClient - from litellm.proxy._experimental.mcp_server.upstream import prepare_mcp_client + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import prepare_mcp_client class CancelledAuth(httpx.Auth): closed = False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 66c5627bc94..979199d0dc9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -10,9 +10,14 @@ This test suite ensures that: """ from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, patch import pytest +from fastapi import HTTPException +from respx import MockRouter + +from litellm.types.mcp import MCPAuth, MCPAuthType from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, @@ -35,6 +40,85 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.get_async_httpx_client" +@pytest.mark.asyncio +@pytest.mark.parametrize("static,forwarded,caller,resolved,expected", [ + ({"Authorization": "Bearer configured"}, {"authorization": "Bearer forwarded"}, None, None, "Bearer configured"), + ({"Authorization": "Bearer configured"}, None, "Bearer caller", None, "Bearer caller"), + ({"Authorization": "Bearer configured"}, None, "Bearer", None, None), + ({"Authorization": "Bearer configured"}, None, "Bearer caller", {"authorization": " "}, None), + ({"Authorization": "Bearer configured"}, None, "Bearer", {"authorization": "Bearer resolved"}, "Bearer resolved"), +]) +async def test_static_auth_validates_headers_after_existing_precedence( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None, + resolved: dict[str, str] | None, expected: str | None, +) -> None: + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.bearer_token, + ) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + caller_token: Final = _request_auth_header.set(caller) + extra_token: Final = _request_extra_headers.set(forwarded) + resolved_token: Final = _request_resolved_auth_headers.set(resolved) + try: + if expected is None: + with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc: + await tool() + assert exc.value.status_code == 500 + assert destination.call_count == 0 + else: + assert await tool() == "authenticated" + assert destination.call_count == 1 + assert destination.calls.last.request.headers["authorization"] == expected + finally: + _request_auth_header.reset(caller_token) + _request_extra_headers.reset(extra_token) + _request_resolved_auth_headers.reset(resolved_token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("credential", ["custom-key", ""]) +async def test_static_auth_uses_configured_custom_header( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, credential: str, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", headers={"x-custom": credential}, + auth_type=MCPAuth.api_key, upstream_token_header="X-Custom", + ) + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + if credential: + assert await tool() == "authenticated" + assert destination.call_count == 1 + assert destination.calls.last.request.headers["x-custom"] == credential + else: + with pytest.raises(HTTPException, match="requires a usable upstream credential"): + await tool() + assert destination.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type,resolved", [ + (MCPAuth.none, None), + (MCPAuth.oauth2, {"Authorization": "Bearer user-oauth"}), +]) +async def test_static_validation_preserves_no_auth_and_resolved_oauth( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthType, resolved: dict[str, str] | None, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function("/echo", "get", {}, "https://upstream.example", auth_type=auth_type) + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="echo") + token: Final = _request_resolved_auth_headers.set(resolved) + try: + assert await tool() == "echo" + assert destination.call_count == 1 + assert destination.calls.last.request.headers.get("authorization") == (resolved or {}).get("Authorization") + finally: + _request_resolved_auth_headers.reset(token) + + def _create_mock_client(method: str, response_text: str, status_code: int = 200) -> AsyncMock: """Utility to create a mocked async httpx client for the given method. From c447c3312db1b99e058062f8cd61f2904fa1e9ef Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 08:13:04 -0700 Subject: [PATCH 089/187] feat(e2e): separate a provider error from a body that failed its rule Build 226's 62 Bedrock rejections are the question this is trying to answer, and "incomplete" would have covered both candidate causes at once. Replaying the completeness rules over eight streams captured from live Bedrock, covering tool use, extended thinking and a max-tokens stop on both streaming endpoints, accepts every one of them, so a rule that is too strict is the less likely half. A provider that answered 429 or 5xx and was retried out of sight is the other, and it now counts as rejected_error_status rather than being folded in with a grammar failure. --- .../code_coverage_tests/test_provider_cache.py | 18 +++++++++++++++--- tests/e2e/PROVIDER_CACHE.md | 2 +- tests/e2e/provider_cache.py | 3 +++ 3 files changed, 19 insertions(+), 4 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 1271b20438a..9668de35a02 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -546,13 +546,25 @@ def test_a_rejection_says_whether_the_body_was_cut_short_or_simply_unfinished( finally: second.shutdown() + refused: Final = cache_edge(store) + provider.status = 429 + provider.response = b'{"message":"Too many requests"}' + third: Final = start_provider_edge(refused, mounts={"openai": upstream}) + try: + call(third.edge.api_base("openai") + "/v1/chat/completions", MARKED) + finally: + third.shutdown() + cut: Final = dict(cut_short.counters.counts) turned_down: Final = dict(unfinished.counters.counts) - assert cut["mount:openai:rejected"] == 1 and turned_down["mount:openai:rejected"] == 1 + errored: Final = dict(refused.counters.counts) + assert cut["mount:openai:rejected"] == turned_down["mount:openai:rejected"] == errored["mount:openai:rejected"] == 1 assert cut["mount:openai:rejected_cut_short"] == 1 - assert "mount:openai:rejected_incomplete" not in cut assert turned_down["mount:openai:rejected_incomplete"] == 1 - assert "mount:openai:rejected_cut_short" not in turned_down + assert errored["mount:openai:rejected_error_status"] == 1 + assert not {"mount:openai:rejected_incomplete", "mount:openai:rejected_error_status"} & set(cut) + assert not {"mount:openai:rejected_cut_short", "mount:openai:rejected_error_status"} & set(turned_down) + assert not {"mount:openai:rejected_cut_short", "mount:openai:rejected_incomplete"} & set(errored) EMBEDDING_SUCCESS: Final = ( diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index 894d9be4efa..ed1e7d517b0 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -54,7 +54,7 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_incomplete` (the body arrived whole and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index f528d08b720..399a0379889 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -55,6 +55,7 @@ EVENTSTREAM_PRELUDE_BYTES: Final = 4 CUT_SHORT: Final = "cut_short" INCOMPLETE: Final = "incomplete" UNREACHABLE: Final = "unreachable" +ERROR_STATUS: Final = "error_status" EVENT_TYPE_HEADER: Final = ":event-type" EVENTSTREAM_HEADERS: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) OPENAI_JSON_PATHS: Final = frozenset({"/v1/chat/completions", "/v1/messages", "/v1/embeddings", "/v1/responses"}) @@ -610,6 +611,8 @@ class CacheEdge: headers: Final = { name: value for name, value in head.headers.items() if name.lower() not in UNRECORDED_RESPONSE_HEADERS } + if not 200 <= head.status_code < 300: + return ERROR_STATUS chunks: Final = capture.chunks() if not successful_response(mount, url, head.status_code, headers, b"".join(chunks)): return INCOMPLETE From 1972a30defcea23d170ba431af99f2e82e652b24 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 08:37:25 -0700 Subject: [PATCH 090/187] revert(e2e): unmount Gemini, its api_base means two things Build 227 mounted Gemini and turned TestGeminiFiles::test_gemini_file_upload red. litellm's two Gemini endpoints disagree about what api_base means. Chat composes {api_base}/models/{model}:{endpoint} and defaults api_base to https://generativelanguage.googleapis.com/v1beta, so the version lives inside it. File upload composes {api_base}/upload/v1beta/files and defaults to the host root, so the version lives outside it. A single api_base cannot satisfy both, and a registration carries no signal about which endpoint the deployment will be used for, so the edge cannot route one and not the other. Backing it out rather than working around it. The cache must never turn a passing test red, which is the same rule the Bedrock model allowlist follows, and Gemini was 7 of roughly 1030 edge calls in that build. Anyone pointing litellm's Gemini provider at an AI gateway or a corporate proxy hits this too, so the fix belongs in litellm; mounting Gemini is one line once it lands. This reverts commit 8a553ceb58887c8aa2aa24c7cfb75ce662d2e321. --- .../test_provider_cache.py | 154 +----------------- tests/e2e/PROVIDER_CACHE.md | 18 +- tests/e2e/provider_cache.py | 46 ------ tests/e2e/provider_cache_routing.py | 3 +- tests/e2e/provider_edge.py | 1 - 5 files changed, 7 insertions(+), 215 deletions(-) diff --git a/tests/code_coverage_tests/test_provider_cache.py b/tests/code_coverage_tests/test_provider_cache.py index 9668de35a02..e57d33a0406 100644 --- a/tests/code_coverage_tests/test_provider_cache.py +++ b/tests/code_coverage_tests/test_provider_cache.py @@ -912,7 +912,7 @@ def test_anthropic_stream_requires_start_finish_and_stop() -> None: assert not successful_response("anthropic", url, 200, headers, start + finish) -@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", ""), ("gemini", "")]) +@pytest.mark.parametrize("provider,suffix", [("openai", "/v1"), ("anthropic", "")]) def test_normal_registration_routes_supported_providers(provider: str, suffix: str) -> None: params: Final = LiteLLMParamsBody(model=f"{provider}/test", api_key="os.environ/SYNTHETIC_KEY", timeout=12) routed: Final = route_cache_model(params, lambda mount: f"http://edge.invalid/{mount}", enabled=True) @@ -924,8 +924,6 @@ def test_normal_registration_routes_supported_providers(provider: str, suffix: s @pytest.mark.parametrize("params", [ LiteLLMParamsBody(model="bedrock/test"), LiteLLMParamsBody(model="azure/test"), - LiteLLMParamsBody(model="vertex_ai/gemini-2.5-flash"), - LiteLLMParamsBody(model="gemini/gemini-2.5-flash", api_base="https://custom.invalid"), LiteLLMParamsBody(model="openai/test", api_base="https://custom.invalid/v1"), LiteLLMParamsBody(model="openai/test", api_base=""), LiteLLMParamsBody(model="openai/test", litellm_credential_name="named-credential"), @@ -1294,153 +1292,3 @@ class TestBedrockStreams: def test_a_stream_that_errored_before_it_started_is_not_recordable(self) -> None: assert not successful_response(BEDROCK_MOUNT, CONVERSE_STREAM_URL, 503, {}, CONVERSE_STREAM_OK) - - -GEMINI_MODEL: Final = "gemini-2.5-flash" -GEMINI_API_VERSION: Final = "/v1beta" -GEMINI_GENERATE_PATH: Final = f"/models/{GEMINI_MODEL}:generateContent" -GEMINI_STREAM_PATH: Final = f"/models/{GEMINI_MODEL}:streamGenerateContent" -GEMINI_USAGE: Final = {"promptTokenCount": 7, "candidatesTokenCount": 1, "totalTokenCount": 25} - - -def gemini_body(finish_reason: str | None, usage: bool = True, candidates: bool = True) -> JsonValue: - candidate: Final[dict[str, JsonValue]] = {"content": {"parts": [{"text": "OK"}], "role": "model"}, "index": 0} - return { - "candidates": [{**candidate, "finishReason": finish_reason} if finish_reason else candidate] - if candidates else [], - **({"usageMetadata": GEMINI_USAGE} if usage else {}), - "modelVersion": GEMINI_MODEL, - } - - -def gemini_unary(finish_reason: str | None = "STOP", usage: bool = True, candidates: bool = True) -> bytes: - return json.dumps(gemini_body(finish_reason, usage, candidates)).encode() - - -def gemini_stream(*finish_reasons: str | None) -> bytes: - return b"".join( - b"data: " + json.dumps(gemini_body(reason)).encode() + b"\r\n\r\n" for reason in finish_reasons - ) - - -@contextmanager -def gemini_edge(cache: CacheEdge, provider: Provider, path: str) -> Generator[str, None, None]: - upstream: Final = f"http://127.0.0.1:{provider.server_port}{GEMINI_API_VERSION}" - running: Final = start_provider_edge(cache, mounts={"gemini": upstream}) - try: - yield running.edge.api_base("gemini") + path - finally: - running.shutdown() - - -class TestGemini: - """Gemini reaches the edge by path prefix alone: litellm composes - `{api_base}/models/{model}:{endpoint}` and sends a static `x-goog-api-key`, - so nothing has to be re-signed and nothing leaves the cache key. The response - grammar is its own though, and the streaming one is the interesting half: every - chunk repeats `usageMetadata`, so only `finishReason` on the last chunk - separates a finished turn from a dropped connection.""" - - @pytest.mark.parametrize("path,response", [ - (GEMINI_GENERATE_PATH, gemini_unary()), - (GEMINI_STREAM_PATH, gemini_stream(None, None, "STOP")), - ], ids=["generate", "stream"]) - def test_a_finished_turn_replays_on_the_next_run( - self, store: RedisResponseStore, provider: Provider, path: str, response: bytes, - ) -> None: - provider.stream = path == GEMINI_STREAM_PATH - provider.response = response - for _ in range(2): - with gemini_edge(cache_edge(store), provider, path) as url: - assert call(url, MARKED).body == response - assert len(provider.hits) == 1 - - @pytest.mark.parametrize("reason", ["MAX_TOKENS", "SAFETY", "RECITATION"]) - def test_a_turn_the_provider_ended_for_its_own_reasons_is_still_finished( - self, store: RedisResponseStore, provider: Provider, reason: str, - ) -> None: - """Reading `finishReason` as a string rather than comparing it to STOP is - deliberate. A turn cut off by the token limit or a safety filter is over, - and rejecting those would send every one of them upstream forever.""" - provider.response = gemini_unary(reason) - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: - assert call(url, MARKED).body == provider.response - assert len(provider.hits) == 1 - - @pytest.mark.parametrize("response", [ - gemini_unary(None), - gemini_unary("STOP", usage=False), - gemini_unary("STOP", candidates=False), - b'{"error":{"code":400,"message":"API key not valid","status":"INVALID_ARGUMENT"}}', - ], ids=["no-finish-reason", "no-usage", "no-candidates", "error-body"]) - def test_an_unfinished_or_failed_turn_never_enters_the_cache( - self, store: RedisResponseStore, provider: Provider, response: bytes, - ) -> None: - provider.response = response - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: - assert call(url, MARKED).body == response - assert len(provider.hits) == 2 - - @pytest.mark.parametrize("response", [ - gemini_stream(None, None), - gemini_stream("STOP", None), - gemini_stream(), - ], ids=["cut-before-the-reason", "reason-then-another-chunk", "empty"]) - def test_a_stream_that_never_named_a_reason_calls_the_provider_every_time( - self, store: RedisResponseStore, provider: Provider, response: bytes, - ) -> None: - provider.stream = True - provider.response = response - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_STREAM_PATH) as url: - assert call(url, MARKED).body == response - assert len(provider.hits) == 2 - - def test_a_response_whose_candidates_did_not_all_finish_is_not_recordable( - self, store: RedisResponseStore, provider: Provider, - ) -> None: - """A request for more than one candidate is answered by more than one, and - the turn is over only when every one of them names a reason. Holding the - whole list to that rule rather than its first entry is what keeps a - half-finished answer from being stored and replayed as a finished one.""" - finished: Final = json.loads(gemini_unary("STOP"))["candidates"][0] - unfinished: Final = json.loads(gemini_unary(None))["candidates"][0] - provider.response = json.dumps( - {"candidates": [finished, {**unfinished, "index": 1}], "usageMetadata": GEMINI_USAGE} - ).encode() - for _ in range(2): - with gemini_edge(cache_edge(store), provider, GEMINI_GENERATE_PATH) as url: - assert call(url, MARKED).body == provider.response - assert len(provider.hits) == 2 - - @pytest.mark.parametrize("path,cacheable", [ - (GEMINI_GENERATE_PATH, True), - (GEMINI_STREAM_PATH, True), - (f"/models/{GEMINI_MODEL}:countTokens", False), - (f"/models/{GEMINI_MODEL}:embedContent", False), - ("/v1/chat/completions", False), - (f"/files/{GEMINI_MODEL}:generateContent", False), - ]) - @pytest.mark.parametrize("version", ["", GEMINI_API_VERSION], ids=["bare", "versioned"]) - def test_only_the_generate_endpoints_are_cacheable(self, version: str, path: str, cacheable: bool) -> None: - """The mount's upstream base carries the API version, so the path the cache - sees is the upstream one and starts `/v1beta`. A rule anchored at the start - of the path would pass every test against a stub with no version prefix and - then cache nothing at all in a real run.""" - assert cacheable_endpoint("gemini", "POST", f"https://gemini.invalid{version}{path}", MARKED) is cacheable - - def test_the_bodies_these_tests_build_match_a_real_gemini_response(self) -> None: - """The shapes above are hand-built so a test can express the turn it means. - This holds them to the fields a live `generativelanguage.googleapis.com` - answer carries, captured 2026-09-16 against gemini-2.5-flash.""" - captured: Final = json.loads( - '{"candidates":[{"content":{"parts":[{"text":"OK"}],"role":"model"},"finishReason":"STOP",' - '"index":0}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":1,' - '"totalTokenCount":25},"modelVersion":"gemini-2.5-flash","responseId":"1J6qauKFI8ut1MkPgNjI4AI"}' - ) - built: Final = json.loads(gemini_unary()) - assert captured.keys() >= built.keys() - assert captured["candidates"][0].keys() >= built["candidates"][0].keys() - assert successful_response("gemini", GEMINI_GENERATE_PATH, 200, {}, json.dumps(captured).encode()) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index ed1e7d517b0..aca81f26c5a 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -1,23 +1,13 @@ # Shared provider-response cache -`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI, Anthropic and Gemini model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live +`E2E_PROVIDER_CACHE=1` enables automatic response reuse in the live E2E mode. Standard OpenAI and Anthropic model registrations use the provider edge, as do Anthropic-on-Bedrock registrations that carry no AWS identity of their own. Existing custom API bases, named credentials, mocked models and realtime WebSocket deployments keep their existing routing. Other provider protocols remain live -The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount, and for `models/{model}:generateContent` and `:streamGenerateContent` on the Gemini mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored +The edge caches complete successful POST responses for `/v1/chat/completions`, `/v1/messages`, `/v1/embeddings` and `/v1/responses` on the OpenAI and Anthropic mounts, SSE streams included, and for `/model/{id}/converse` and `/model/{id}/invoke` on a Bedrock mount. Unsupported endpoints pass through. Each endpoint family has its own completeness rule, so a truncated embedding or a Responses run that never reached `response.completed` is not stored Bedrock's streaming endpoints, `converse-stream` and `invoke-with-response-stream`, cache too. AWS frames those as binary `vnd.amazon.eventstream` rather than SSE, so botocore's own parser reads the frames and validates both CRCs, and each endpoint is then held to its terminal grammar. That matters more than the endpoint count suggests: the Claude Code compat cells drive the real CLI, which always streams, so streaming is most of the suite's Bedrock traffic Two details of that rule are worth knowing before changing it. A ConverseStream ends with `metadata`, not with `messageStop`, and the `metadata` frame is what carries the token usage litellm prices the call from, so the rule requires it: a stream cut between the two still names a stop reason but would replay as a free call. And a dropped connection is invisible to the parser, which yields the frames it did receive and silently discards a trailing partial one, so the body is also checked against the frame lengths it declares. A stream cut one byte short parses clean and has to be caught that way -## Gemini - -Gemini needs nothing that Bedrock needed. litellm composes `{api_base}/models/{model}:{endpoint}` from a custom api_base, so a path-prefixed mount reaches it, and the credential travels as a static `x-goog-api-key` header that no host rewrite invalidates. Nothing is re-signed and nothing is excluded from the key, so a recording still cannot cross credentials - -The mount's upstream base carries the API version, which is the one detail worth remembering: the path the cache rules see is the upstream one, `/v1beta/models/...`, not the one the proxy sent. A rule anchored at the start of that path would look right against a local stub and then cache nothing at all in a real run - -A finished turn names a `finishReason` on every candidate and reports `usageMetadata`. The reason is read as a string rather than compared to `STOP`, because `MAX_TOKENS` and the safety reasons end a turn just as finally and rejecting them would send every one of them upstream forever. Streaming is the more interesting half: Gemini repeats `usageMetadata` on every chunk and names a `finishReason` only on the last one, so the terminator is the final event rather than any event, and a stream the connection cut short ends on a chunk with usage and no reason - -Vertex is not mounted. litellm grafts the default Vertex path onto an api_base only when that api_base has no path of its own, so a Vertex mount needs a root-mounted edge on its own port rather than a path prefix. Gemini and Vertex are separate providers in litellm and the Gemini mount does not cover Vertex deployments - ## Request identity A recording belongs to one test. The key is a keyed digest over the test's node id, the method, the URL, the effective outbound headers (including authentication and HTTP-library defaults), body presence and the body bytes, with one normalization: a 12-hex-digit run, the shape `unique_marker()` mints, is replaced by a placeholder in both the URL and a UTF-8 body. Nothing else is normalized away. No prompts, JSON values or credentials are rewritten, and the rule is the one `fixture_canonical.py` already applies for record/replay, so there is a single definition of what a marker is @@ -40,7 +30,9 @@ Only deployments that carry no AWS identity of their own route to the edge. A de Which models route is an explicit allowlist in `provider_cache_routing.py`, mirroring the runner role's IAM policy, which names its models one by one. That coupling is deliberate: the edge re-signs with the run pod's identity, so a model the role cannot invoke comes back 403 from Bedrock rather than falling back. An unlisted model keeps its direct path and loses only caching, so adding a Bedrock model to the suite can never turn it red. Adding one to the edge is a policy edit in litellm-ops plus a line here -Vertex is not mounted. litellm's `_check_custom_proxy` rewrites a path-prefixed Vertex `api_base` into `{api_base}:{endpoint}`, dropping project, location and model, so a mount under a path prefix cannot work without either a root-mounted edge on its own port or a change in litellm. Gemini is a separate provider there and does have a working path-prefixed form, so it is mounted; see the Gemini section +Vertex and Gemini are not mounted, for different reasons. litellm grafts the default Vertex path onto an `api_base` only when that `api_base` has no path of its own, so a path-prefixed Vertex mount instead becomes `{api_base}:{endpoint}`, dropping project, location and model. Vertex needs a root-mounted edge on its own port, or a change in litellm + +Gemini reaches a path-prefixed mount perfectly well and was mounted for one build, then backed out, because litellm's two Gemini endpoints disagree about what `api_base` means. Chat composes `{api_base}/models/{model}:{endpoint}` and defaults `api_base` to `https://generativelanguage.googleapis.com/v1beta`, so the version has to be inside it. File upload composes `{api_base}/upload/v1beta/files` and defaults to the host root, so the version has to be outside it. One `api_base` cannot satisfy both, and a deployment gives no signal at registration time about which it will be used for, so mounting Gemini turned `TestGeminiFiles::test_gemini_file_upload` red in build 227. Anyone pointing litellm's Gemini provider at an AI gateway or a corporate proxy hits the same thing; it is a litellm bug rather than a cache limitation, and mounting Gemini is one line once it is fixed Recordings are shared across workers and builds through dedicated Redis, separate from the candidate's own cache. They expire 86,400 seconds after capture starts, based on Redis time. Reads never extend expiry. There is no scheduled recapture: the next miss calls the provider again. Bounded coordination reduces duplicate concurrent calls, but slow or failed captures may lead to extra live calls after the wait expires diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 399a0379889..444972ffce5 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -37,10 +37,6 @@ SIGNATURE_HEADERS: Final = frozenset( {"authorization", "x-amz-date", "x-amz-security-token", "x-amz-content-sha256"} ) BEDROCK_MOUNT_PREFIX: Final = "bedrock" -GEMINI_MOUNT: Final = "gemini" -GEMINI_MODELS_SEGMENT: Final = "/models" -GEMINI_GENERATE_SUFFIX: Final = ":generateContent" -GEMINI_STREAM_SUFFIX: Final = ":streamGenerateContent" BEDROCK_CONVERSE_SUFFIX: Final = "/converse" BEDROCK_INVOKE_SUFFIX: Final = "/invoke" BEDROCK_CONVERSE_STREAM_SUFFIX: Final = "/converse-stream" @@ -162,21 +158,12 @@ def is_bedrock(mount: str) -> bool: return mount.partition("/")[0] == BEDROCK_MOUNT_PREFIX -def is_gemini(mount: str) -> bool: - return mount == GEMINI_MOUNT - - def cacheable_endpoint(mount: str, method: str, url: str, body: bytes | None) -> bool: if method != "POST" or body is None or len(body) > MAX_REQUEST_BYTES: return False path: Final = urlsplit(url).path if is_bedrock(mount): return path.startswith("/model/") and path.endswith(BEDROCK_SUFFIXES) - if is_gemini(mount): - collection, _, resource = path.rpartition("/") - return collection.endswith(GEMINI_MODELS_SEGMENT) and resource.endswith( - (GEMINI_GENERATE_SUFFIX, GEMINI_STREAM_SUFFIX) - ) return path in OPENAI_JSON_PATHS @@ -203,8 +190,6 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, for value in values ): return False - if is_gemini(mount): - return complete_gemini_stream(values) if urlsplit(url).path == "/v1/responses": return complete_responses_stream(values) if urlsplit(url).path == "/v1/chat/completions": @@ -216,8 +201,6 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, return False if not isinstance(value, dict) or value.get("error") is not None: return False - if is_gemini(mount): - return complete_gemini_candidates(value) path: Final = urlsplit(url).path if path == "/v1/messages": return value.get("type") == "message" and isinstance(value.get("content"), list) and isinstance(value.get("stop_reason"), str) @@ -236,35 +219,6 @@ def successful_response(mount: str, url: str, status: int, headers: Mapping[str, ) -def complete_gemini_candidates(value: Mapping[str, JsonValue]) -> bool: - """A finished Gemini turn names a ``finishReason`` on every candidate and - reports the usage litellm prices the call from. ``finishReason`` is read as a - string rather than compared to ``STOP`` because ``MAX_TOKENS`` and the safety - reasons end a turn just as finally, and a cache that rejected them would send - every one of them upstream forever.""" - candidates: Final = value.get("candidates") - return ( - isinstance(value.get("usageMetadata"), dict) - and isinstance(candidates, list) - and bool(candidates) - and all( - isinstance(candidate, dict) and isinstance(candidate.get("finishReason"), str) - for candidate in candidates - ) - ) - - -def complete_gemini_stream(values: tuple[JsonValue, ...]) -> bool: - """Gemini repeats ``usageMetadata`` on every chunk but names a - ``finishReason`` only on the last one, so the terminator is the final event - rather than any event. A stream the connection cut short ends on a chunk that - carries usage and no reason, which is exactly what this rejects.""" - if not values: - return False - last: Final = values[-1] - return isinstance(last, dict) and complete_gemini_candidates(last) - - def complete_bedrock_response(url: str, body: bytes) -> bool: """Converse answers with ``output`` plus a ``stopReason``; InvokeModel on an Anthropic model answers the Anthropic message shape. Either way a truncated diff --git a/tests/e2e/provider_cache_routing.py b/tests/e2e/provider_cache_routing.py index c4b02beac2d..f9775a2b152 100644 --- a/tests/e2e/provider_cache_routing.py +++ b/tests/e2e/provider_cache_routing.py @@ -19,7 +19,6 @@ BEDROCK_EDGE_MODELS: Final = frozenset( } ) ENV_REFERENCE_PREFIX: Final = "os.environ/" -EDGE_PROVIDERS: Final = frozenset({"openai", "anthropic", "gemini"}) def bedrock_region(declared: str | None) -> str: @@ -82,7 +81,7 @@ def route_cache_model( return route_bedrock(params, base_for, mode) if mode == "realtime" or params.api_base is not None: return params - if provider not in EDGE_PROVIDERS: + if provider not in {"openai", "anthropic"}: return params base: Final = base_for(provider) if base is None: diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 219df55233a..2606b26fe99 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -104,7 +104,6 @@ EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( { "openai": "https://api.openai.com", "anthropic": "https://api.anthropic.com", - "gemini": "https://generativelanguage.googleapis.com/v1beta", **{ f"bedrock/{region}": f"https://bedrock-runtime.{region}.amazonaws.com" for region in BEDROCK_REGIONS From 6c517bfc49eea5535fd76bfc66df62d239bd94d0 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:00:17 -0700 Subject: [PATCH 091/187] fix(mcp): reject scheme-only API key authorization payloads --- .../outbound_credentials/adapter.py | 5 ++++ .../outbound_credentials/test_adapter.py | 23 ++++++++++++++- .../mcp_server/test_mcp_server_manager.py | 5 +++- .../test_openapi_to_mcp_generator.py | 29 +++++++++++++++++++ 4 files changed, 60 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 5358878a248..fb9af25933a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -401,6 +401,11 @@ def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> b return True if value.lower() in ("bearer", "basic", "token", "apikey"): return False + if auth_type == MCPAuth.api_key: + api_scheme: Final = value.split(None, 1)[0] + if api_scheme.lower() in ("bearer", "token", "apikey"): + api_credential: Final = strip_auth_scheme(value, api_scheme).strip() + return api_credential.lower() != api_scheme.lower() if auth_type in (MCPAuth.bearer_token, MCPAuth.token): scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token" credential: Final = strip_auth_scheme(value, scheme).strip() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 2885fdaef95..4ada7c1763d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -7,6 +7,7 @@ maps each CredError onto its HTTP status. These pin the parity-critical mapping import base64 from types import SimpleNamespace +from typing import Final import pytest from fastapi import HTTPException @@ -20,7 +21,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import raise_user_oauth_challenge, to_server_spec, to_subject, + validate_static_credential, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, @@ -34,10 +37,28 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( SharedKey, TokenExchangeConfig, ) -from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer +@pytest.mark.parametrize("auth_type,header,value", [ + (MCPAuth.api_key, "Authorization", "Bearer fixture-key"), + (MCPAuth.api_key, "Authorization", "ApiKey fixture-key"), + (MCPAuth.api_key, "Authorization", "token fixture-key"), + (MCPAuth.api_key, "Authorization", "Bearer token"), + (MCPAuth.api_key, "Authorization", "opaque-key"), + (MCPAuth.api_key, "Authorization", "Custom Custom"), + (MCPAuth.api_key, "X-API-Key", "Bearer Bearer"), + (MCPAuth.api_key, "X-Custom", "ApiKey ApiKey"), + (MCPAuth.authorization, "Authorization", "Bearer Bearer"), +]) +def test_static_credential_preserves_supported_api_key_and_raw_headers( + auth_type: MCPAuthType, header: str, value: str, +) -> None: + result: Final = validate_static_credential(auth_type, {header: value}, upstream_token_header=header) + assert isinstance(result, Ok) + + def _server(**kwargs) -> MCPServer: return MCPServer(server_id="s", name="n", transport=MCPTransport.http, **kwargs) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 4c4c45162ca..9601f4dac4f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13767,7 +13767,10 @@ class TestProtectedCredentialPreparation: assert custom_slot is None or custom_slot not in request.headers @pytest.mark.asyncio - @pytest.mark.parametrize("value", ["", " ", "Bearer", "Basic", "token", "ApiKey"]) + @pytest.mark.parametrize("value", [ + "", " ", "Bearer", "Basic", "token", "ApiKey", + "Bearer Bearer", "ApiKey ApiKey", "token token", "bEaReR BEARER", "aPiKeY\tAPIKEY", + ]) async def test_api_key_rejects_authorization_without_a_credential(self, value: str) -> None: server: Final = MCPServer( server_id="caller-empty", name="caller-empty", url="https://upstream.example/mcp", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 979199d0dc9..a7b7b0e9b44 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -40,6 +40,35 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.get_async_httpx_client" +@pytest.mark.asyncio +@pytest.mark.parametrize("value,accepted", [ + ("Bearer Bearer", False), ("ApiKey ApiKey", False), ("token token", False), + ("bEaReR BEARER", False), ("aPiKeY\tAPIKEY", False), + ("Bearer fixture-key", True), ("ApiKey fixture-key", True), ("token fixture-key", True), +]) +async def test_api_key_authorization_validates_payload_before_http( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, value: str, accepted: bool, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", auth_type=MCPAuth.api_key, + ) + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + caller_token: Final = _request_auth_header.set(value) + try: + if accepted: + assert await tool() == "authenticated" + assert destination.call_count == 1 + assert destination.calls.last.request.headers["authorization"] == value + else: + with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc: + await tool() + assert exc.value.status_code == 500 + assert destination.call_count == 0 + finally: + _request_auth_header.reset(caller_token) + + @pytest.mark.asyncio @pytest.mark.parametrize("static,forwarded,caller,resolved,expected", [ ({"Authorization": "Bearer configured"}, {"authorization": "Bearer forwarded"}, None, None, "Bearer configured"), From 7d42bc751debd3ffa7eaa6166d20f627cee0f067 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 10:05:13 -0700 Subject: [PATCH 092/187] docs(e2e): name all four rejection reasons in the counter note --- tests/e2e/PROVIDER_CACHE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index aca81f26c5a..e06b3c01653 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -46,7 +46,7 @@ The trusted runner receives: - `E2E_PROVIDER_CACHE_NAMESPACE`: shared environment namespace, independent of build and candidate revision - `E2E_PROVIDER_CACHE_METRICS_DIR`: optional per-process counter artifact directory -Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all one or the other is a different problem, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits +Do not give cache credentials to candidate deployments. Counter artifacts contain no recorded payloads or credentials. Hits count shared-cache responses; upstream attempts count actual forwards from the edge. A rejection also counts its reason, one of `rejected_cut_short` (the consumer walked away mid-capture), `rejected_error_status` (the provider answered, with an error), `rejected_incomplete` (the body arrived whole with a success status and failed its endpoint's rule) or `rejected_unreachable` (the provider could not be reached at all). A mount whose rejections are nearly all of one kind is a different problem from one whose rejections are nearly all of another, and the flat count cannot tell them apart. Every counter is emitted twice, once as a flat total and once under `mount:{mount}:`, so a hit rate can be read per provider rather than only in aggregate. Existing application-cache observations still count requests arriving at the edge, including shared-cache hits Tests that require real provider timing, limits or state use `@pytest.mark.provider_live`. The marker keeps newly registered models on live routes without weakening their assertions. The provider prompt-caching tests carry it because a replayed priming response reports cache creation rather than a cache read. From 847172d311f625eba25c97b7c70cf317a735175c Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 17:31:04 +0000 Subject: [PATCH 093/187] fix(anthropic): carry the served model from message_start onto stream chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/chat/handler.py | 4 + .../chat/test_anthropic_chat_handler.py | 78 +++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 4dd0deeb62b..e3912138d8c 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -632,6 +632,7 @@ class ModelResponseIterator: self.tool_name_reverse_map: dict[str, str] = tool_name_reverse_map or {} # Generate response ID once per stream to match OpenAI-compatible behavior self.response_id = _generate_id() + self.served_model: str | None = None # Track if we're currently streaming a response_format tool self.is_response_format_tool: bool = False @@ -1067,6 +1068,8 @@ class ModelResponseIterator: } """ message_start_block: Final = MessageStartBlock(**chunk) + start_message: Final = message_start_block["message"] + self.served_model = start_message["model"] if "model" in start_message else None if "usage" in message_start_block["message"]: usage = self._handle_usage(anthropic_usage_chunk=message_start_block["message"]["usage"]) elif type_chunk == "error": @@ -1098,6 +1101,7 @@ class ModelResponseIterator: ], usage=usage, id=self.response_id, + model=self.served_model, ) return returned_chunk diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index c3400dc40c3..1e6639ef537 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2719,3 +2719,81 @@ class TestRustChatCompletionsHook: "model": "m", "messages": [], } + + +def _served_model_stream_chunks(model: str | None) -> list[dict]: + message: Final = { + "id": "msg_served", + "type": "message", + "role": "assistant", + "content": [], + "usage": {"input_tokens": 10, "output_tokens": 1}, + } + if model is not None: + message["model"] = model + return [ + {"type": "message_start", "message": message}, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 2}, + }, + {"type": "message_stop"}, + ] + + +def test_message_start_model_is_carried_on_stream_chunks(): + iterator: Final = ModelResponseIterator(None, sync_stream=True) + + parsed: Final = [ + iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks("claude-served-1") + ] + + assert all(chunk.model == "claude-served-1" for chunk in parsed) + + +def test_message_start_without_model_leaves_chunk_model_unset(): + iterator: Final = ModelResponseIterator(None, sync_stream=True) + + parsed: Final = [ + iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks(None) + ] + + assert all(chunk.model is None for chunk in parsed) + + +def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper(): + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + served_model: Final = "claude-served-1" + sse_lines: Final = [ + f"data: {json.dumps(chunk)}\n".encode() for chunk in _served_model_stream_chunks(served_model) + ] + iterator: Final = ModelResponseIterator(iter(sse_lines), sync_stream=True) + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(iterator), + model="anthropic/claude-requested", + custom_llm_provider="anthropic", + logging_obj=MagicMock(), + ) + + chunks: Final = list(wrapper) + + assert len(chunks) > 1 + for chunk in chunks[1:]: + assert chunk._hidden_params["provider_response_model"] == served_model + assembled: Final = litellm.stream_chunk_builder( + chunks=list(chunks), messages=[{"role": "user", "content": "hi"}] + ) + assert assembled._hidden_params["provider_response_model"] == served_model From 70ef8b24b6675faa14e7e3ec006d79debe032609 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:44:00 -0700 Subject: [PATCH 094/187] fix(mcp): reject bare schemes in raw authorization --- .../outbound_credentials/adapter.py | 2 +- .../outbound_credentials/test_adapter.py | 2 +- .../mcp_server/test_mcp_server_manager.py | 35 ++++++++++++++++--- .../test_openapi_to_mcp_generator.py | 20 +++++++---- 4 files changed, 46 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index fb9af25933a..ba223f73b2d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -397,7 +397,7 @@ _STATIC_MODES: Final = frozenset( def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool: if not value: return False - if auth_type == MCPAuth.authorization or (auth_type == MCPAuth.api_key and name != "authorization"): + if auth_type == MCPAuth.api_key and name != "authorization": return True if value.lower() in ("bearer", "basic", "token", "apikey"): return False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 4ada7c1763d..78da9ff4d77 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -50,7 +50,7 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer (MCPAuth.api_key, "Authorization", "Custom Custom"), (MCPAuth.api_key, "X-API-Key", "Bearer Bearer"), (MCPAuth.api_key, "X-Custom", "ApiKey ApiKey"), - (MCPAuth.authorization, "Authorization", "Bearer Bearer"), + (MCPAuth.authorization, "Authorization", "opaque-secret-value"), ]) def test_static_credential_preserves_supported_api_key_and_raw_headers( auth_type: MCPAuthType, header: str, value: str, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 9601f4dac4f..186c46e1b37 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -13577,19 +13577,46 @@ class TestProtectedCredentialPreparation: assert exc.value.status_code in (401, 500) @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type,slot", [(MCPAuth.api_key, "X-API-Key"), (MCPAuth.authorization, "Authorization")]) - async def test_raw_static_value_named_token_is_a_usable_credential(self, auth_type: MCPAuthType, slot: str) -> None: + @pytest.mark.parametrize("auth_type,slot,value", [ + (MCPAuth.api_key, "X-API-Key", "token"), + (MCPAuth.authorization, "Authorization", "opaque-secret-value"), + (MCPAuth.authorization, "Authorization", "Bearer abc"), + (MCPAuth.authorization, "Authorization", "Custom abc"), + ]) + async def test_raw_static_credentials_are_forwarded_unchanged( + self, auth_type: MCPAuthType, slot: str, value: str, + ) -> None: server = MCPServer(server_id="raw-key", name="raw-key", url="https://upstream.example/mcp", - transport=MCPTransport.http, auth_type=auth_type, authentication_token="token") + transport=MCPTransport.http, auth_type=auth_type, authentication_token=value) client = await MCPServerManager()._create_mcp_client(server) assert client._resolved_auth is not None request = httpx.Request("GET", server.url) flow = client._resolved_auth.auth_flow(request) try: - assert next(flow).headers[slot] == "token" + assert next(flow).headers[slot] == value finally: flow.close() + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["Bearer", "basic", "token", "ApiKey", " bEaReR ", "\tTOKEN\t"]) + @pytest.mark.parametrize("source", ["configured", "caller", "forwarded"]) + async def test_raw_authorization_rejects_bare_schemes_before_dispatch( + self, respx_mock: MockRouter, value: str, source: str, + ) -> None: + server: Final = MCPServer( + server_id="raw-empty", name="raw-empty", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.authorization, + authentication_token=value if source == "configured" else None, + ) + destination: Final = respx_mock.route().respond(200) + with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc: + await MCPServerManager()._create_mcp_client( + server, mcp_auth_header=value if source == "caller" else None, + extra_headers={"Authorization": value} if source == "forwarded" else None, + ) + assert exc.value.status_code == 500 + assert destination.call_count == 0 + @pytest.mark.asyncio async def test_byok_flag_cannot_bypass_incomplete_obo(self) -> None: server = MCPServer(server_id="obo-byok", name="obo-byok", url="https://upstream.example/mcp", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index a7b7b0e9b44..a9def20e75d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -41,17 +41,23 @@ GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp @pytest.mark.asyncio -@pytest.mark.parametrize("value,accepted", [ - ("Bearer Bearer", False), ("ApiKey ApiKey", False), ("token token", False), - ("bEaReR BEARER", False), ("aPiKeY\tAPIKEY", False), - ("Bearer fixture-key", True), ("ApiKey fixture-key", True), ("token fixture-key", True), +@pytest.mark.parametrize("auth_type,value,accepted", [ + (MCPAuth.api_key, "Bearer Bearer", False), (MCPAuth.api_key, "ApiKey ApiKey", False), + (MCPAuth.api_key, "token token", False), (MCPAuth.api_key, "bEaReR BEARER", False), + (MCPAuth.api_key, "aPiKeY\tAPIKEY", False), (MCPAuth.api_key, "Bearer fixture-key", True), + (MCPAuth.api_key, "ApiKey fixture-key", True), (MCPAuth.api_key, "token fixture-key", True), + (MCPAuth.authorization, "Bearer", False), (MCPAuth.authorization, "basic", False), + (MCPAuth.authorization, "token", False), (MCPAuth.authorization, "ApiKey", False), + (MCPAuth.authorization, " bEaReR ", False), (MCPAuth.authorization, "\tTOKEN\t", False), + (MCPAuth.authorization, "opaque-secret-value", True), (MCPAuth.authorization, "Bearer abc", True), + (MCPAuth.authorization, "Custom abc", True), ]) -async def test_api_key_authorization_validates_payload_before_http( - respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, value: str, accepted: bool, +async def test_authorization_validates_credentials_before_http( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, auth_type: MCPAuthType, value: str, accepted: bool, ) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") tool: Final = create_tool_function( - "/echo", "get", {}, "https://upstream.example", auth_type=MCPAuth.api_key, + "/echo", "get", {}, "https://upstream.example", auth_type=auth_type, ) destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") caller_token: Final = _request_auth_header.set(value) From 07d49364285b083190393893600101e186158a59 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 17:44:59 +0000 Subject: [PATCH 095/187] fix(anthropic): satisfy strict lint and update tests pinned to the dropped served model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/anthropic/chat/handler.py | 3 ++- .../chat/test_anthropic_chat_handler.py | 16 ++++------------ .../guardrail_hooks/test_tool_permission.py | 7 ++++--- ...test_anthropic_passthrough_logging_handler.py | 6 +++--- 4 files changed, 13 insertions(+), 19 deletions(-) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index e3912138d8c..58ed9b38ccb 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -1069,7 +1069,8 @@ class ModelResponseIterator: """ message_start_block: Final = MessageStartBlock(**chunk) start_message: Final = message_start_block["message"] - self.served_model = start_message["model"] if "model" in start_message else None + if "model" in start_message: + self.served_model = start_message["model"] if "usage" in message_start_block["message"]: usage = self._handle_usage(anthropic_usage_chunk=message_start_block["message"]["usage"]) elif type_chunk == "error": diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index 1e6639ef537..7db44b4709e 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2756,9 +2756,7 @@ def _served_model_stream_chunks(model: str | None) -> list[dict]: def test_message_start_model_is_carried_on_stream_chunks(): iterator: Final = ModelResponseIterator(None, sync_stream=True) - parsed: Final = [ - iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks("claude-served-1") - ] + parsed: Final = [iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks("claude-served-1")] assert all(chunk.model == "claude-served-1" for chunk in parsed) @@ -2766,9 +2764,7 @@ def test_message_start_model_is_carried_on_stream_chunks(): def test_message_start_without_model_leaves_chunk_model_unset(): iterator: Final = ModelResponseIterator(None, sync_stream=True) - parsed: Final = [ - iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks(None) - ] + parsed: Final = [iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks(None)] assert all(chunk.model is None for chunk in parsed) @@ -2777,9 +2773,7 @@ def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper(): from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper served_model: Final = "claude-served-1" - sse_lines: Final = [ - f"data: {json.dumps(chunk)}\n".encode() for chunk in _served_model_stream_chunks(served_model) - ] + sse_lines: Final = [f"data: {json.dumps(chunk)}\n".encode() for chunk in _served_model_stream_chunks(served_model)] iterator: Final = ModelResponseIterator(iter(sse_lines), sync_stream=True) wrapper: Final = CustomStreamWrapper( completion_stream=iter(iterator), @@ -2793,7 +2787,5 @@ def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper(): assert len(chunks) > 1 for chunk in chunks[1:]: assert chunk._hidden_params["provider_response_model"] == served_model - assembled: Final = litellm.stream_chunk_builder( - chunks=list(chunks), messages=[{"role": "user", "content": "hi"}] - ) + assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=[{"role": "user", "content": "hi"}]) assert assembled._hidden_params["provider_response_model"] == served_model diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 87a1b84acc5..427fa43ffd5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -1291,8 +1291,9 @@ class TestToolPermissionGuardrailAnthropicMessages: async def test_rewrite_mode_keeps_the_stream_identity_it_had_before_the_shared_helper(self): """Well-formed SSE must round-trip exactly as it did before the helpers were shared. - The shared module can stamp the upstream message id and model onto the assembled response - for callers that ask for it; this path never did, and a client reads those bytes. + The shared module can stamp the upstream message id onto the assembled response for + callers that ask for it; this path never did, and a client reads those bytes. The model, + though, is now the upstream's, matching what the untouched passthrough shows clients. """ with patch.object(self.rewriting, "should_run_guardrail", return_value=True): out = await self._drain(self.rewriting, self._sse_chunks("Read")) @@ -1304,7 +1305,7 @@ class TestToolPermissionGuardrailAnthropicMessages: if line.startswith("data: ") and json.loads(line[6:]).get("type") == "message_start" )["message"] assert message_start["id"].startswith("chatcmpl-"), "the rewritten stream must not adopt the upstream message id" - assert message_start["model"] == "unknown-model", "the rewritten stream must not adopt the upstream model" + assert message_start["model"] == "claude-sonnet-4-5", "the rewritten stream reports the model the upstream served" @pytest.mark.asyncio async def test_message_start_without_a_dict_message_fails_closed(self): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 2d7397594aa..ba8b5fa3ac4 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2541,7 +2541,7 @@ class TestRecordPartialUsageForFailure: function_id="test-partial-usage-failure", ) - def _interrupted_chunks(self): + def _interrupted_chunks(self, *, model: str = "claude-sonnet-5"): return [ self._sse( "message_start", @@ -2551,7 +2551,7 @@ class TestRecordPartialUsageForFailure: "id": "msg_abc", "type": "message", "role": "assistant", - "model": "claude-sonnet-5", + "model": model, "content": [], "stop_reason": None, "stop_sequence": None, @@ -2588,7 +2588,7 @@ class TestRecordPartialUsageForFailure: AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure( litellm_logging_obj=logging_obj, request_body={"model": "claude-unpriced-test-model", "stream": True}, - all_chunks=self._interrupted_chunks(), + all_chunks=self._interrupted_chunks(model="claude-unpriced-test-model"), ) usage = logging_obj.model_call_details["combined_usage_object"] From 686556fffc475a0ae083286440b50bba771da26c Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 17:49:15 +0000 Subject: [PATCH 096/187] test(anthropic): build served-model stream chunks immutably Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/test_anthropic_chat_handler.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index 7db44b4709e..a241fc03d57 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2721,18 +2721,19 @@ class TestRustChatCompletionsHook: } -def _served_model_stream_chunks(model: str | None) -> list[dict]: - message: Final = { - "id": "msg_served", - "type": "message", - "role": "assistant", - "content": [], - "usage": {"input_tokens": 10, "output_tokens": 1}, - } - if model is not None: - message["model"] = model +def _served_model_stream_chunks(model: str | None) -> list[dict[str, object]]: return [ - {"type": "message_start", "message": message}, + { + "type": "message_start", + "message": { + "id": "msg_served", + "type": "message", + "role": "assistant", + "content": [], + "usage": {"input_tokens": 10, "output_tokens": 1}, + **({"model": model} if model is not None else {}), + }, + }, { "type": "content_block_start", "index": 0, From ecd72c51da019f30debf182d44f89e2c2b6e4854 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 11:29:25 -0700 Subject: [PATCH 097/187] refactor(ui): register ssh skill sources verbatim and share the https host rules The ssh parser rebuilt the clone url it was given, stripping a trailing .git and appending one back. Git treats that suffix as optional, and the hosts whose clone paths are not org/repo break when it is forced on, so an Azure DevOps v3 or a CodeCommit v1/repos url registered through the form would no longer clone. It also carried its own host pattern, which demanded an alphabetic final label and so rejected internal hosts like gitlab.internal.k8s2 that the https path accepts. Rewrite the scp form into an ssh:// url purely to validate it, reuse the https host and credential checks through a shared isSafeHost, and store exactly what the user typed. Only a url that survives the round trip unchanged is accepted, which is what keeps traversal segments out of the feed, so the two ssh regexes, the dots-only guard and the clone-url builders all collapse into one function. --- .../claude_code_plugins/helpers.test.ts | 36 +++++--- .../components/claude_code_plugins/helpers.ts | 84 +++++++++---------- 2 files changed, 65 insertions(+), 55 deletions(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts index ccdc0a1a388..5bf782aabf3 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts @@ -476,21 +476,30 @@ describe("parseSkillSource", () => { source: "url", url: "git@ghe.example.com:org/repo.git", }); - expect(parseSkillSource("git@ghe.example.com:org/repo")?.parsed).toEqual({ - source: "url", - url: "git@ghe.example.com:org/repo.git", - }); expect(parseSkillSource("git@ghe.example.com:org/repo.git")?.suggestedName).toBe("repo"); }); - it("normalizes an ssh:// clone url and keeps a custom port", () => { - expect(parseSkillSource("ssh://git@ghe.example.com/org/repo")?.parsed).toEqual({ + it("stores an ssh clone url exactly as typed, so a forced .git suffix cannot break azure devops or codecommit", () => { + for (const url of [ + "git@ghe.example.com:org/repo", + "git@ssh.dev.azure.com:v3/org/project/repo", + "ssh://git@ghe.example.com/org/repo", + "ssh://apka1234@git-codecommit.us-east-1.amazonaws.com/v1/repos/my-repo", + "ssh://git@ghe.example.com:2222/org/nested/repo.git", + ]) { + expect(parseSkillSource(url)?.parsed).toEqual({ source: "url", url }); + } + expect(parseSkillSource("git@ssh.dev.azure.com:v3/org/project/repo")?.suggestedName).toBe("repo"); + }); + + it("accepts an internal host whose last label is not alphabetic, matching the https rule", () => { + expect(parseSkillSource("git@gitlab.internal.k8s2:org/repo.git")?.parsed).toEqual({ source: "url", - url: "ssh://git@ghe.example.com/org/repo.git", + url: "git@gitlab.internal.k8s2:org/repo.git", }); - expect(parseSkillSource("ssh://git@ghe.example.com:2222/org/nested/repo.git")?.parsed).toEqual({ + expect(parseSkillSource("https://gitlab.internal.k8s2/org/repo")?.parsed).toEqual({ source: "url", - url: "ssh://git@ghe.example.com:2222/org/nested/repo.git", + url: "https://gitlab.internal.k8s2/org/repo", }); }); @@ -510,17 +519,22 @@ describe("parseSkillSource", () => { expect(parseSkillSource("ssh://ghe.example.com/org/repo.git")).toBeNull(); }); - it("rejects ssh remotes with ip hosts or dot-only path segments", () => { + it("rejects ssh remotes with ip hosts or traversal segments", () => { expect(parseSkillSource("git@10.0.0.5:org/repo.git")).toBeNull(); expect(parseSkillSource("ssh://git@169.254.169.254/org/repo")).toBeNull(); expect(parseSkillSource("git@ghe.example.com:../etc")).toBeNull(); expect(parseSkillSource("ssh://git@ghe.example.com/org/../repo")).toBeNull(); + expect(parseSkillSource("git@ghe.example.com:org/../../etc/passwd")).toBeNull(); expect(parseSkillSource("git@ghe.example.com:org/.github")?.parsed).toEqual({ source: "url", - url: "git@ghe.example.com:org/.github.git", + url: "git@ghe.example.com:org/.github", }); }); + it("rejects an ssh remote carrying a password, which would publish a secret on the feed", () => { + expect(parseSkillSource("ssh://git:s3cret@ghe.example.com/org/repo.git")).toBeNull(); + }); + it("returns null for empty and garbage input", () => { expect(parseSkillSource("")).toBeNull(); expect(parseSkillSource(" ")).toBeNull(); diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts index 3c11a6cedbc..b4d9eab08ec 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts @@ -29,22 +29,33 @@ export const SHA256_REGEX = /^[0-9a-fA-F]{64}$/; export const isValidSha256 = (digest: string): boolean => digest.trim() === "" || SHA256_REGEX.test(digest.trim()); -// WHATWG normalizes obfuscated IPv4 (e.g. 2130706433, 0x7f.0.0.1) to dotted-decimal, so this -// catches every IPv4 form; bracketed IPv6 is rejected separately. +// WHATWG normalizes obfuscated IPv4 (e.g. 2130706433, 0x7f.0.0.1) to dotted-decimal on https, so +// this catches every IPv4 form there; on a non-special scheme like ssh it catches the dotted form +// only. Bracketed IPv6 is rejected separately. const IPV4_HOST_REGEX = /^\d{1,3}(\.\d{1,3}){3}$/; const GITHUB_ORG_REGEX = /^[A-Za-z0-9-]+$/; const GITHUB_REPO_REGEX = /^[A-Za-z0-9._-]+$/; -const SSH_SCP_REGEX = /^([a-z0-9._-]+)@([a-z0-9.-]+\.[a-z]{2,}):([a-z0-9._-]+(?:\/[a-z0-9._-]+)+?)(?:\.git)?\/?$/i; -const SSH_URL_REGEX = - /^ssh:\/\/([a-z0-9._-]+)@([a-z0-9.-]+\.[a-z]{2,})(:\d+)?\/([a-z0-9._-]+(?:\/[a-z0-9._-]+)+?)(?:\.git)?\/?$/i; -const DOTS_ONLY_SEGMENT_REGEX = /^\.+$/; +const SSH_SCHEME = "ssh://"; +const SSH_SCP_REGEX = /^([a-z0-9._-]+)@([^:/@]+):(?!\/)(.+)$/i; const buildRepoUrl = (url: URL): string => `${url.protocol}//${url.host}${url.pathname.replace(/\/+$/, "")}`; const pathSegments = (url: URL): string[] => url.pathname.split("/").filter((seg) => seg !== ""); +const toUrl = (candidate: string): URL | null => { + try { + return new URL(candidate); + } catch { + return null; + } +}; + +/** One host rule for every scheme, so an ssh remote is neither more nor less trusted than its https twin. */ +const isSafeHost = (url: URL): boolean => + url.hostname.includes(".") && !url.hostname.startsWith("[") && !IPV4_HOST_REGEX.test(url.hostname); + /** * Validate and normalize a repository URL into a parsed URL, or null. Enforces https (rejects * http/ssh/git/etc.), rejects embedded credentials, and requires a dotted host, so the public @@ -57,20 +68,8 @@ const parseRepoUrl = (raw: string): URL | null => { return null; } const withScheme = /^[a-z][a-z0-9+.-]*:\/\//i.test(trimmed) ? trimmed : `https://${trimmed}`; - let url: URL; - try { - url = new URL(withScheme); - } catch { - return null; - } - if ( - url.protocol !== "https:" || - url.username !== "" || - url.password !== "" || - !url.hostname.includes(".") || - url.hostname.startsWith("[") || - IPV4_HOST_REGEX.test(url.hostname) - ) { + const url = toUrl(withScheme); + if (!url || url.protocol !== "https:" || url.username !== "" || url.password !== "" || !isSafeHost(url)) { return null; } return url; @@ -178,32 +177,29 @@ const parseRawGitSource = (url: URL, subPath?: string): SkillSourcePreview | nul return buildGitSourcePreview("Git", buildRepoUrl(url), repoName, subPath); }; -interface SshRemote { - cloneUrl: string; - repoName: string; -} - -const buildSshRemote = (rawPath: string, toCloneUrl: (repoPath: string) => string): SshRemote | null => { - if (rawPath.split("/").some((segment) => DOTS_ONLY_SEGMENT_REGEX.test(segment))) { +/** + * Parse an scp-style `git@host:org/repo` or `ssh://git@host/org/repo` clone URL, registering it + * exactly as typed: git treats the `.git` suffix as optional, and forcing one on breaks hosts whose + * paths are not `org/repo`, like Azure DevOps `v3/...` and CodeCommit `v1/repos/...`. The scp form is + * rewritten to `ssh://` only to reuse the https host and credential rules, and only a URL that + * survives that round trip unchanged is accepted, which keeps traversal segments off the feed. + */ +const parseSshSource = (raw: string, subPath?: string): SkillSourcePreview | null => { + const trimmed = raw.trim(); + const scp = SSH_SCP_REGEX.exec(trimmed); + const candidate = scp ? `${SSH_SCHEME}${scp[1]}@${scp[2]}/${scp[3]}` : trimmed; + if (!candidate.toLowerCase().startsWith(SSH_SCHEME)) { return null; } - const bare = rawPath.replace(/\.git$/i, ""); - return { cloneUrl: toCloneUrl(`${bare}.git`), repoName: lastSegment(bare) }; -}; - -const parseSshRemote = (raw: string): SshRemote | null => { - const trimmed = raw.trim(); - const sshUrl = SSH_URL_REGEX.exec(trimmed); - if (sshUrl) { - const [, user, host, port, path] = sshUrl; - return buildSshRemote(path, (repoPath) => `ssh://${user}@${host}${port ?? ""}/${repoPath}`); + const url = toUrl(candidate); + if (!url || url.username === "" || url.password !== "" || !isSafeHost(url)) { + return null; } - const scp = SSH_SCP_REGEX.exec(trimmed); - if (scp) { - const [, user, host, path] = scp; - return buildSshRemote(path, (repoPath) => `${user}@${host}:${repoPath}`); + const pathStart = candidate.indexOf("/", SSH_SCHEME.length); + if (pathStart === -1 || url.pathname !== candidate.slice(pathStart) || pathSegments(url).length < 2) { + return null; } - return null; + return buildGitSourcePreview("SSH", trimmed, lastSegment(url.pathname).replace(/\.git$/i, ""), subPath); }; const parseArchiveSource = (url: URL): SkillSourcePreview => ({ @@ -220,9 +216,9 @@ const parseArchiveSource = (url: URL): SkillSourcePreview => ({ * with an optional subfolder turning it into git-subdir. */ export const parseSkillSource = (rawUrl: string, subPath?: string): SkillSourcePreview | null => { - const ssh = parseSshRemote(rawUrl); + const ssh = parseSshSource(rawUrl, subPath); if (ssh) { - return buildGitSourcePreview("SSH", ssh.cloneUrl, ssh.repoName, subPath); + return ssh; } const url = parseRepoUrl(rawUrl); if (!url) { From 0143fe5583e845af20606c9d6becfbff35b9c8c4 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 18:30:12 +0000 Subject: [PATCH 098/187] fix(guardrails): hold tool-call windows until the final scan and expose Bedrock streaming flags to the UI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/guardrail_translation/handler.py | 4 +- .../guardrail_translation/base_translation.py | 6 +- .../chat/guardrail_translation/handler.py | 4 +- .../guardrail_translation/handler.py | 15 ++- .../proxy/guardrails/guardrail_endpoints.py | 6 +- .../guardrail_hooks/bedrock_guardrails.py | 6 +- .../unified_guardrail/unified_guardrail.py | 14 +++ .../test_anthropic_guardrail_handler.py | 3 + .../test_openai_guardrail_handler.py | 3 + ...test_openai_responses_guardrail_handler.py | 21 +++++ .../test_bedrock_guardrails.py | 38 ++++++++ .../test_streaming_buffer_until_moderated.py | 91 ++++++++++++++++++- .../guardrails/test_guardrail_endpoints.py | 16 ++++ 13 files changed, 217 insertions(+), 10 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 2ea20143f0c..93d87585496 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1561,10 +1561,12 @@ class AnthropicMessagesHandler(BaseTranslation): def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: stream_ended: Final = self._check_streaming_has_ended(responses_so_far) + tool_use_fingerprints: Final = self._streamed_tool_use_fingerprints(responses_so_far) return StreamingScanKey( texts=(self.get_streaming_string_so_far(responses_so_far),), - tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (), + tool_calls=tool_use_fingerprints if stream_ended else (), stream_ended=stream_ended, + tool_calls_in_flight=bool(tool_use_fingerprints) and not stream_ended, ) @classmethod diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index f1143425ced..89ad67f0485 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -40,11 +40,15 @@ class StreamingScanKey: """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys compare equal when the round would scan the same content again; ``stream_ended`` stays out of the comparison and only says whether the handler is on its - end-of-stream path, where an empty payload is still scanned today.""" + end-of-stream path, where an empty payload is still scanned today. + ``tool_calls_in_flight`` also stays out of the comparison: it flags that tool + calls have streamed which this round cannot scan yet, so a buffered window + holding them must stay withheld until the end-of-stream scan covers them.""" texts: tuple[str, ...] tool_calls: tuple[str, ...] = () stream_ended: bool = field(default=False, compare=False) + tool_calls_in_flight: bool = field(default=False, compare=False) @property def has_nothing_to_scan(self) -> bool: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 01e14f2248d..9b818696bbb 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -792,10 +792,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream)) stream_ended: Final = self._first_choice_has_finished(responses_so_far) + tool_call_fingerprints: Final = self._streamed_tool_call_fingerprints(responses_so_far) return StreamingScanKey( texts=tuple(self._combine_streaming_texts(chunks).values()), - tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (), + tool_calls=tool_call_fingerprints if stream_ended else (), stream_ended=stream_ended, + tool_calls_in_flight=bool(tool_call_fingerprints) and not stream_ended, ) @staticmethod diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 27ff55f120c..321f0a0d6d3 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1177,9 +1177,22 @@ class OpenAIResponsesHandler(BaseTranslation): return None if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value: return self._completed_response_scan_key(stream_item_field(last_event, "response")) + stream_ended: Final = self._check_streaming_has_ended(responses_so_far) return StreamingScanKey( texts=(self.get_streaming_string_so_far(responses_so_far),), - stream_ended=self._check_streaming_has_ended(responses_so_far), + stream_ended=stream_ended, + tool_calls_in_flight=not stream_ended and self._has_streamed_tool_call_events(responses_so_far), + ) + + @staticmethod + def _has_streamed_tool_call_events(responses_so_far: Sequence[object]) -> bool: + return any( + stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES + or ( + stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES + and stream_item_field(stream_item_field(event, "item"), "type") in _TOOL_CALL_ITEM_TYPES + ) + for event in responses_so_far ) @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index afb9997f2e6..6874d7aa73e 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -37,6 +37,7 @@ from litellm.types.guardrails import ( ApplyGuardrailResponse, BaseLitellmParams, BedrockGuardrailConfigModel, + BedrockGuardrailStreamingParams, Guardrail, GuardrailEventHooks, GuardrailInfoResponse, @@ -1959,7 +1960,10 @@ async def get_provider_specific_params(): ``` """ # Get fields from the models - bedrock_fields: Final = _get_fields_from_model(BedrockGuardrailConfigModel) + bedrock_fields: Final = { + **_get_fields_from_model(BedrockGuardrailConfigModel), + **_get_fields_from_model(BedrockGuardrailStreamingParams), + } presidio_fields: Final = _get_fields_from_model(PresidioPresidioConfigModelUserInterface) lakera_v2_fields: Final = _get_fields_from_model(LakeraV2GuardrailConfigModel) tool_permission_fields: Final = _get_fields_from_model(ToolPermissionGuardrailConfigModel) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index db68c7b5c3f..434c52ca6f3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -330,7 +330,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra)) def _streams_incrementally(self) -> bool: - return not self.streaming_buffer_until_moderated and not self.mask_response_content + if self.mask_response_content: + return False + if not self.streaming_buffer_until_moderated: + return True + return self.streaming_buffer_release_on_scan and not self.streaming_end_of_stream_only @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 70dc6eec24c..5ac9070faa6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -1032,6 +1032,7 @@ class UnifiedLLMGuardrails(CustomLogger): # message (True) vs emit a standalone block message (False, buffered). chunks_yielded = False last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round + tool_calls_in_flight = False # rebind-ok: tracks the latest scan key's unscanned tool calls async for item in response: chunk_counter += 1 @@ -1079,6 +1080,9 @@ class UnifiedLLMGuardrails(CustomLogger): if chunk_counter % sampling_rate == 0: endpoint_translation = mappings[CallTypes(call_type)]() scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far) + if scan_key is not None: + tool_calls_in_flight = scan_key.tool_calls_in_flight + hold_window = buffer_until_moderated and tool_calls_in_flight if _is_redundant_scan(scan_key, last_scan_key): verbose_proxy_logger.debug( "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round", @@ -1086,6 +1090,8 @@ class UnifiedLLMGuardrails(CustomLogger): guardrail_to_apply.guardrail_name, ) if buffer_until_moderated: + if hold_window: + continue for withheld_item in withheld_items: chunks_yielded = True responses_yielded.append(withheld_item) @@ -1151,6 +1157,14 @@ class UnifiedLLMGuardrails(CustomLogger): return if scan_key is not None: last_scan_key = scan_key + if hold_window: + verbose_proxy_logger.debug( + "Holding %s buffered chunks for guardrail %s: streamed tool calls await the end-of-stream scan", + len(withheld_items), + guardrail_to_apply.guardrail_name, + ) + withheld_items[:] = original_items + continue for original_item in original_items: chunks_yielded = True responses_yielded.append(original_item) 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 7522e9a62e5..cb18a5f192c 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 @@ -2482,7 +2482,10 @@ class TestAnthropicMessagesHandlerStreamingScanKey: open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use]) ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")]) assert open_key == StreamingScanKey(texts=("hi",)) + assert open_key.tool_calls_in_flight is True + assert handler.get_streaming_scan_key([self._text_delta("hi")]).tool_calls_in_flight is False assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key.tool_calls_in_flight is False assert ended_key != open_key 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 cb884fb7cc1..0e59e1d1590 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 @@ -2206,8 +2206,11 @@ class TestStreamingScanKey: [self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")] ) assert open_key == StreamingScanKey(texts=("hi",)) + assert open_key.tool_calls_in_flight is True + assert handler.get_streaming_scan_key([self._chunk("hi")]).tool_calls_in_flight is False assert ended_key.texts == ("hi",) assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key.tool_calls_in_flight is False assert ended_key != open_key def test_text_after_the_first_choice_finishes_still_changes_the_key(self): diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 48d86384633..60375b03404 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -3211,3 +3211,24 @@ class TestOpenAIResponsesHandlerStreamingScanKey: def test_output_item_done_round_is_never_deduped(self): done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}} assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None + + def test_streamed_tool_call_events_flag_tool_calls_in_flight_until_the_stream_ends(self): + handler = OpenAIResponsesHandler() + added = { + "type": "response.output_item.added", + "sequence_number": 1, + "item": {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "get_weather"}, + } + arguments_delta = { + "type": "response.function_call_arguments.delta", + "sequence_number": 2, + "item_id": "fc_1", + "delta": '{"city":', + } + function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"} + assert handler.get_streaming_scan_key([self._delta(0, "hi")]).tool_calls_in_flight is False + assert handler.get_streaming_scan_key([self._delta(0, "hi"), added]).tool_calls_in_flight is True + assert handler.get_streaming_scan_key([self._delta(0, "hi"), arguments_delta]).tool_calls_in_flight is True + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), added, self._completed(3, [function_call])]) + assert ended_key.tool_calls_in_flight is False + assert len(ended_key.tool_calls) == 1 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 dc275c912c8..1d3d7a452b6 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 @@ -5724,6 +5724,44 @@ async def test_buffered_default_hook_scans_before_any_chunk(): assert len([e for e in events if e != "scan"]) >= 1 +@pytest.mark.asyncio +async def test_buffered_release_on_scan_hook_releases_each_window_after_its_scan(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-release-on-scan", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + streaming_buffer_release_on_scan=True, + streaming_sampling_rate=1, + ) + + assert guardrail._streams_incrementally() is True + events = await _run_streaming_hook_recording_order(guardrail) + + assert events == ["scan", ("chunk", "Hello"), "scan", ("chunk", " world"), ("chunk", "")] + + +@pytest.mark.asyncio +async def test_buffered_release_on_scan_defers_to_end_of_stream_only(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-release-on-scan-end-only", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + streaming_buffer_release_on_scan=True, + streaming_end_of_stream_only=True, + streaming_sampling_rate=1, + ) + + assert guardrail._streams_incrementally() is False + events = await _run_streaming_hook_recording_order(guardrail) + + assert events.count("scan") == 1 + assert events[0] == "scan" + + @pytest.mark.asyncio async def test_masking_keeps_buffered_path_even_when_unbuffered_configured(): guardrail = BedrockGuardrail( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py index 89d51866533..cbc17e680f0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py @@ -23,10 +23,18 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) -from litellm.types.utils import Delta, GenericGuardrailAPIInputs, ModelResponseStream, StreamingChoices +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + GenericGuardrailAPIInputs, + ModelResponseStream, + StreamingChoices, +) BLOCK_MESSAGE = "Blocked by policy: this response was withheld." ORIGINAL_MARKER = "ORIGINAL-SECRET-ANSWER" +TOOL_ARGUMENTS_MARKER = "TOOL-ARGS-SECRET" class _BlockingGuardrail(CustomGuardrail): @@ -76,6 +84,24 @@ class _CountingPassingGuardrail(_PassingGuardrail): return inputs +class _ToolCallRecordingGuardrail(_CountingPassingGuardrail): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.tool_call_scan_indexes: List[int] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.scan_count += 1 + if inputs.get("tool_calls"): + self.tool_call_scan_indexes.append(self.scan_count) + return inputs + + class _SecondScanBlockingGuardrail(_CountingPassingGuardrail): async def apply_guardrail( self, @@ -165,20 +191,63 @@ def _chat_chunk(content: str = "", finish_reason: str | None = None) -> ModelRes ) +def _tool_call_chunk(arguments: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-windowed", + created=1724900000, + model="gpt-4", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_1", + type="function", + index=0, + function=Function(name="run_shell", arguments=arguments), + ) + ], + ), + finish_reason=finish_reason, + ) + ], + ) + + async def _windowed_chat_stream( - yielded_count: List[int], collected: List[Any], content_chunks: List[str] + yielded_count: List[int], + collected: List[Any], + content_chunks: List[str], + tool_argument_chunks: List[str] | None = None, ) -> AsyncGenerator[ModelResponseStream, None]: for content in content_chunks: yielded_count.append(len(collected)) yield _chat_chunk(content) + for arguments in tool_argument_chunks or []: + yielded_count.append(len(collected)) + yield _tool_call_chunk(arguments) yielded_count.append(len(collected)) - yield _chat_chunk(finish_reason="stop") + yield _chat_chunk(finish_reason="tool_calls" if tool_argument_chunks else "stop") + + +def _tool_argument_text(chunks: List[Any]) -> str: + return "".join( + tool_call.function.arguments or "" + for chunk in chunks + if isinstance(chunk, ModelResponseStream) + for choice in chunk.choices + for tool_call in choice.delta.tool_calls or [] + ) async def _run_windowed( guardrail: CustomGuardrail, content_chunks: List[str], end_of_stream_only: bool = False, + tool_argument_chunks: List[str] | None = None, ) -> tuple[List[Any], List[int]]: guardrail.streaming_buffer_until_moderated = True guardrail.streaming_buffer_release_on_scan = True @@ -195,7 +264,7 @@ async def _run_windowed( yielded_count: List[int] = [] async for chunk in unified.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, - response=_windowed_chat_stream(yielded_count, collected, content_chunks), + response=_windowed_chat_stream(yielded_count, collected, content_chunks, tool_argument_chunks), request_data=request_data, ): collected.append(chunk) @@ -282,6 +351,20 @@ async def test_windowed_buffer_drops_blocked_window(): assert '"error"' not in raw +@pytest.mark.asyncio +async def test_windowed_buffer_holds_tool_call_windows_until_end_of_stream_scan(): + guardrail = _ToolCallRecordingGuardrail(guardrail_name="windowed-tools", event_hook="post_call") + content_chunks = ["one ", "two ", "three "] + tool_argument_chunks = ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'] + + collected, yielded_count = await _run_windowed(guardrail, content_chunks, tool_argument_chunks=tool_argument_chunks) + + assert yielded_count == [0, 0, 2, 2, 2, 2, 2] + assert _chat_text(collected) == "".join(content_chunks) + assert _tool_argument_text(collected) == "".join(tool_argument_chunks) + assert guardrail.tool_call_scan_indexes == [guardrail.scan_count] + + @pytest.mark.asyncio async def test_windowed_buffer_with_explicit_end_of_stream_only_stays_fully_buffered(): guardrail = _CountingPassingGuardrail(guardrail_name="windowed-eos", event_hook="post_call") diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 530f8ffd854..bf641fd6cd0 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -682,6 +682,22 @@ async def test_provider_specific_params_includes_embedding_toggle(): assert field["default_value"] is False +@pytest.mark.asyncio +async def test_provider_specific_params_exposes_bedrock_streaming_flags(): + from litellm.proxy.guardrails.guardrail_endpoints import get_provider_specific_params + + provider_params = await get_provider_specific_params() + + bedrock = provider_params["bedrock"] + assert "guardrailIdentifier" in bedrock + assert "guardrailVersion" in bedrock + assert bedrock["streaming_buffer_release_on_scan"]["type"] == "boolean" + assert bedrock["streaming_buffer_release_on_scan"]["default_value"] is False + assert bedrock["streaming_buffer_until_moderated"]["default_value"] is True + assert bedrock["streaming_end_of_stream_only"]["type"] == "boolean" + assert bedrock["streaming_sampling_rate"]["type"] == "number" + + @pytest.mark.asyncio async def test_provider_specific_params_includes_hide_secrets(): """hide-secrets lives in the enterprise package so it is not in From 7cd6869cfaf7f4c84995a8b878800ac73eac65aa Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 18:40:10 +0000 Subject: [PATCH 099/187] test(ui): match skill source links by exact name so codeql stops flagging the host regexes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../components/claude_code_plugins/skill_detail.test.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx index 2f600397a4b..1ec1f72f4a1 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx @@ -15,7 +15,7 @@ const buildSkill = (source: Plugin["source"]): Plugin => ({ describe("SkillDetail source", () => { it("links a github source to the repository", () => { render(); - expect(screen.getByRole("link", { name: /github.com\/org\/repo/ })).toHaveAttribute( + expect(screen.getByRole("link", { name: "github.com/org/repo" })).toHaveAttribute( "href", "https://github.com/org/repo", ); @@ -26,7 +26,7 @@ describe("SkillDetail source", () => { , ); expect(screen.getByText("git@ghe.example.com:org/repo.git")).toBeInTheDocument(); - expect(screen.queryByRole("link", { name: /ghe.example.com/ })).not.toBeInTheDocument(); + expect(screen.queryByRole("link")).not.toBeInTheDocument(); }); it("renders an ssh git-subdir source as plain text without a tree path", () => { @@ -37,6 +37,6 @@ describe("SkillDetail source", () => { />, ); expect(screen.getByText("git@ghe.example.com:org/repo.git @ plugins/x")).toBeInTheDocument(); - expect(screen.queryByRole("link", { name: /ghe.example.com/ })).not.toBeInTheDocument(); + expect(screen.queryByRole("link")).not.toBeInTheDocument(); }); }); From 6e5cdd8f121be08fc94eac6ff5b0cf69c9c7ce79 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 18:45:01 +0000 Subject: [PATCH 100/187] fix(router): count TPM/RPM usage before building rate-limit headers Router.make_call now increments the deployment TPM/RPM counter before set_response_headers reads remaining usage, so the headers carry post-increment values directly and the in-flight subtraction workaround from LIT-2719 is removed. deployment_callback_on_success adds only the tokens not yet counted (streams) and never a second request. The counter key uses the resolved deployment name so wildcard routes are read back correctly, and the proxy strips the router-owned counted-tokens marker from client metadata Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/litellm_pre_call_utils.py | 9 +- litellm/router.py | 236 +++++++++--------- .../add_retry_fallback_headers.py | 9 +- .../track_deployment_metrics.py | 23 ++ .../test_router_helper_utils.py | 217 +++++++++++----- .../proxy/test_litellm_pre_call_utils.py | 8 +- 7 files changed, 321 insertions(+), 182 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index ce5b65080ee..f9b9da621e1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1499,6 +1499,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags" +ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 563db811edc..dce9a496622 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -25,6 +25,7 @@ from litellm.constants import ( LITELLM_PROXY_MASTER_KEY_ALIAS, OTEL_SERVICE_NAME_METADATA_KEYS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, @@ -336,7 +337,13 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg # and read by spend logs as fact; a client value has no legitimate meaning and no # key or team setting keeps it, so the strip is never gated. _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( - {"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY} + { + "attempted_fallbacks", + "original_model_group", + "request_retry_count", + CLIENT_OUTPUT_CEILING_METADATA_KEY, + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, + } ) _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override" diff --git a/litellm/router.py b/litellm/router.py index 789fc81d8d3..f8abd9be758 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -63,6 +63,7 @@ from litellm.constants import ( DEFAULT_MAX_LRU_CACHE_SIZE, INTERNAL_CALL_ORIGIN_METADATA_KEY, OUTPUT_TOKEN_CEILING_PARAMS, + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, RUNTIME_UPDATABLE_ROUTER_SETTINGS, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, @@ -132,7 +133,7 @@ from litellm.router_utils.add_retry_fallback_headers import ( get_hidden_params_dict, prepare_response_for_header_attachment, replace_complexity_router_headers, - response_in_flight_token_count, + response_total_token_count, ) from litellm.router_utils.auto_router_model_naming import ( AUTO_ROUTER_MODEL_PREFIX, @@ -215,6 +216,8 @@ from litellm.router_utils.reasoning_effort_capability import ( resolve_supported_reasoning_efforts, ) from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + find_deployment_metadata, + get_counted_usage_tokens, increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, ) @@ -7937,6 +7940,7 @@ class Router: response = original_function(*args, **kwargs) if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): response = await response + await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) ## PROCESS RESPONSE HEADERS response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs) @@ -8153,8 +8157,6 @@ class Router: """ Track remaining tpm/rpm quota for model in model_list """ - from litellm.types.caching import RedisPipelineIncrementOperation - try: # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): @@ -8162,114 +8164,132 @@ class Router: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: raise ValueError("standard_logging_object is None") - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - deployment_name: Final = kwargs["litellm_params"]["metadata"].get( - "deployment", None - ) # stable name - works for wildcard routes as well - # Get model_group and id from kwargs like the sync version does - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - model_info: Final = kwargs["litellm_params"].get("model_info", {}) or {} - id = model_info.get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) + litellm_params: Final = kwargs["litellm_params"] + metadata: Final = litellm_params.get("metadata") + if metadata is None: + return + model_group: Final = metadata.get("model_group", None) + model_info: Final = litellm_params.get("model_info", {}) or {} + deployment_id: Final = model_info.get("id", None) + if model_group is None or deployment_id is None or self.get_deployment(model_id=str(deployment_id)) is None: + return - ## get deployment info - deployment_info: Final = self.get_deployment(model_id=id) + # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits + increment_deployment_successes_for_current_minute( + litellm_router_instance=self, + deployment_id=str(deployment_id), + ) - if deployment_info is None: - return - else: - deployment_model_info: Final = self.get_router_model_info( - deployment=deployment_info, - received_model_name=model_group, - ) - # get tpm/rpm from deployment info - tpm: Final = deployment_info.get("tpm", None) - rpm: Final = deployment_info.get("rpm", None) - - ## check tpm/rpm in litellm_params - tpm_litellm_params: Final = deployment_info.litellm_params.tpm - rpm_litellm_params: Final = deployment_info.litellm_params.rpm - - ## check tpm/rpm in model_info - tpm_model_info: Final = deployment_model_info.get("tpm", None) - rpm_model_info: Final = deployment_model_info.get("rpm", None) - - # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits - increment_deployment_successes_for_current_minute( - litellm_router_instance=self, - deployment_id=id, - ) - - deployment_dict = deployment_info if isinstance(deployment_info, dict) else deployment_info.model_dump() - has_io_token_limits: Final = deployment_has_io_token_limits(deployment_dict) - - ## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are - ## set. IO deployments still record TPM/RPM usage here so TPM-aware - ## routing strategies see their real load in mixed model groups; their - ## itpm/otpm enforcement runs separately in ModelRateLimitingCheck. - if ( - tpm is None - and rpm is None - and tpm_litellm_params is None - and rpm_litellm_params is None - and tpm_model_info is None - and rpm_model_info is None - and not has_io_token_limits - ): - return - - parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) - total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0) - - # ------------ - # Setup values - # ------------ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - - tpm_key = RouterCacheEnum.TPM.value.format(id=id, current_minute=current_minute, model=deployment_name) - # ------------ - # Update usage - # ------------ - # update cache - pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [] - - ## TPM - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=tpm_key, - increment_value=total_tokens, - ttl=RoutingArgs.ttl.value, - ) - ) - - ## RPM - rpm_key = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name) - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=rpm_key, - increment_value=1, - ttl=RoutingArgs.ttl.value, - ) - ) - - await self.cache.async_increment_cache_pipeline( - increment_list=pipeline_operations, - parent_otel_span=parent_otel_span, - ) - - return tpm_key + total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0) + counted_tokens: Final = get_counted_usage_tokens(litellm_params) + deployment_name: Final = metadata.get("deployment", None) + return await self._increment_deployment_usage( + deployment_id=str(deployment_id), + deployment_name=deployment_name if isinstance(deployment_name, str) else None, + model_group=model_group, + total_tokens=total_tokens if counted_tokens is None else max(0, total_tokens - counted_tokens), + rpm_increment=1 if counted_tokens is None else 0, + parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), + ) except Exception as e: verbose_router_logger.debug( "litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e ) + async def increment_deployment_usage_for_response( + self, + response: object, + request_kwargs: dict[str, object], + ) -> None: + """Count the request before the headers are read; the success callback adds only what is still missing""" + if response is None: + return + try: + deployment_metadata: Final = find_deployment_metadata(request_kwargs) + model_group: Final = request_kwargs.get("model") + if deployment_metadata is None or not isinstance(model_group, str): + return + model_info: Final = deployment_metadata["model_info"] + deployment_id: Final = model_info.get("id") if isinstance(model_info, dict) else None + if deployment_id is None: + return + total_tokens: Final = response_total_token_count(response) + deployment_name: Final = deployment_metadata.get("deployment") + deployment_metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] = total_tokens + try: + await self._increment_deployment_usage( + deployment_id=str(deployment_id), + deployment_name=deployment_name if isinstance(deployment_name, str) else None, + model_group=model_group, + total_tokens=total_tokens, + rpm_increment=1, + parent_otel_span=_get_parent_otel_span_from_kwargs(request_kwargs), + ) + except Exception: + deployment_metadata.pop(ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, None) + raise + except Exception as e: + verbose_router_logger.debug( + "litellm.router.Router::increment_deployment_usage_for_response(): Exception occured - %s", e + ) + + async def _increment_deployment_usage( + self, + *, + deployment_id: str, + deployment_name: str | None, + model_group: str, + total_tokens: float, + rpm_increment: int, + parent_otel_span: Span | None, + ) -> str | None: + from litellm.types.caching import RedisPipelineIncrementOperation + + deployment_info: Final = self.get_deployment(model_id=deployment_id) + if deployment_info is None: + return None + deployment_model_info: Final = self.get_router_model_info( + deployment=deployment_info, + received_model_name=model_group, + ) + configured_limits: Final = ( + deployment_info.get("tpm", None), + deployment_info.get("rpm", None), + deployment_info.litellm_params.tpm, + deployment_info.litellm_params.rpm, + deployment_model_info.get("tpm", None), + deployment_model_info.get("rpm", None), + ) + ## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are + ## set. IO deployments still record TPM/RPM usage here so TPM-aware + ## routing strategies see their real load in mixed model groups; their + ## itpm/otpm enforcement runs separately in ModelRateLimitingCheck. + if all(limit is None for limit in configured_limits) and not deployment_has_io_token_limits( + deployment_info.model_dump() + ): + return None + if total_tokens <= 0 and rpm_increment <= 0: + return None + + current_minute: Final = get_utc_datetime().strftime("%H-%M") # use the same timezone regardless of system clock + tpm_key: Final = RouterCacheEnum.TPM.value.format( + id=deployment_id, current_minute=current_minute, model=deployment_name + ) + rpm_key: Final = RouterCacheEnum.RPM.value.format( + id=deployment_id, current_minute=current_minute, model=deployment_name + ) + pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [ + RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value) + for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment)) + if increment_value > 0 + ] + await self.cache.async_increment_cache_pipeline( + increment_list=pipeline_operations, + parent_otel_span=parent_otel_span, + ) + return tpm_key + def sync_deployment_callback_on_success( self, kwargs, # kwargs to completion @@ -11205,15 +11225,7 @@ class Router: if model_group is not None: remaining_usage: Final = await self.get_remaining_model_group_usage(model_group) - # get_remaining_model_group_usage reads the router's TPM/RPM counter, - # which is incremented post-response by deployment_callback_on_success. - # Replay the in-flight increment for TPM/RPM only (LIT-2719); ITPM/OTPM - # counters are incremented at reservation time and must not be adjusted. - apply_remaining_usage_headers( - additional_headers, - remaining_usage, - response_in_flight_token_count(response), - ) + apply_remaining_usage_headers(additional_headers, remaining_usage) return response def _build_model_name_index(self, model_list: list) -> None: diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index bc88feef7d2..cbca5880b52 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -151,7 +151,7 @@ def apply_quality_router_decision_headers( additional_headers[header] = str(decision[field]) -def response_in_flight_token_count(response: object) -> int: +def response_total_token_count(response: object) -> int: usage: Final = response.get("usage") if isinstance(response, dict) else getattr(response, "usage", None) if usage is None: return 0 @@ -166,15 +166,10 @@ def response_in_flight_token_count(response: object) -> int: def apply_remaining_usage_headers( additional_headers: dict[str, object], remaining_usage: dict[str, int], - in_flight_tokens: int, ) -> None: - in_flight_delta: Final = { - "x-ratelimit-remaining-tokens": in_flight_tokens, - "x-ratelimit-remaining-requests": 1, - } for header, value in remaining_usage.items(): if value is not None and header not in additional_headers: - additional_headers[header] = value - in_flight_delta.get(header, 0) + additional_headers[header] = value def _normalize_hidden_params(hidden_params: object) -> dict[str, object]: diff --git a/litellm/router_utils/router_callbacks/track_deployment_metrics.py b/litellm/router_utils/router_callbacks/track_deployment_metrics.py index 01893e925bc..6b422e98ec8 100644 --- a/litellm/router_utils/router_callbacks/track_deployment_metrics.py +++ b/litellm/router_utils/router_callbacks/track_deployment_metrics.py @@ -9,8 +9,11 @@ get_deployment_failures_for_current_minute get_deployment_successes_for_current_minute """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final +from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY + if TYPE_CHECKING: from litellm.router import Router as _Router @@ -18,6 +21,26 @@ if TYPE_CHECKING: else: LitellmRouter = Any +_METADATA_CHANNELS: Final = ("litellm_metadata", "metadata") + + +def find_deployment_metadata(kwargs: Mapping[str, object]) -> dict[str, object] | None: + buckets: Final = (kwargs.get(channel) for channel in _METADATA_CHANNELS) + return next((bucket for bucket in buckets if isinstance(bucket, dict) and "model_info" in bucket), None) + + +def get_counted_usage_tokens(litellm_params: Mapping[str, object]) -> int | None: + buckets: Final = (litellm_params.get(channel) for channel in _METADATA_CHANNELS) + counted: Final = next( + ( + bucket[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] + for bucket in buckets + if isinstance(bucket, dict) and ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY in bucket + ), + None, + ) + return counted if isinstance(counted, int) and not isinstance(counted, bool) else None + def increment_deployment_successes_for_current_minute( litellm_router_instance: LitellmRouter, diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 14d86743557..215101735c8 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1,3 +1,4 @@ +import asyncio import json import os import traceback @@ -12,7 +13,7 @@ from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload from litellm.types.utils import StandardLoggingPayload from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo -from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS +from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY @pytest.fixture @@ -928,17 +929,10 @@ async def test_set_response_headers(model_list): @pytest.mark.asyncio -async def test_set_response_headers_subtracts_in_flight_delta(model_list): +async def test_set_response_headers_passes_through_post_increment_counters(model_list): """ - LIT-2719: router-derived `x-ratelimit-remaining-*` headers must be - post-decrement (match OpenAI/Anthropic vendor semantics) so the proxy's - HTTP response headers and the prometheus gauges that read them stay - comparable across providers. - - Router's TPM/RPM counter is incremented post-response by - `deployment_callback_on_success`, so `get_remaining_model_group_usage` - sees pre-decrement values. `set_response_headers` must replay the - in-flight increment before writing the headers. + LIT-3058: `make_call` increments the router's TPM/RPM counter before the headers + are built, so `set_response_headers` writes the remaining values it reads as-is. """ from pydantic import BaseModel @@ -952,49 +946,10 @@ async def test_set_response_headers_subtracts_in_flight_delta(model_list): router = Router(model_list=model_list) router.get_remaining_model_group_usage = AsyncMock( return_value={ - "x-ratelimit-remaining-tokens": 1000, + "x-ratelimit-remaining-tokens": 958, "x-ratelimit-limit-tokens": 1000, - "x-ratelimit-remaining-requests": 100, + "x-ratelimit-remaining-requests": 99, "x-ratelimit-limit-requests": 100, - } - ) - - resp = _Resp() - resp._hidden_params = {} - await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") - - headers = resp._hidden_params["additional_headers"] - assert headers["x-ratelimit-remaining-tokens"] == 958 - assert headers["x-ratelimit-remaining-requests"] == 99 - # Limit headers pass through unmodified. - assert headers["x-ratelimit-limit-tokens"] == 1000 - assert headers["x-ratelimit-limit-requests"] == 100 - - -@pytest.mark.asyncio -async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_list): - """ - The in-flight replay applies only to the post-incremented TPM/RPM counters - (`x-ratelimit-remaining-tokens` / `-requests`). The ITPM/OTPM counters are - incremented at reservation time (pre-call), so the input/output token - headers already reflect this request and must pass through untouched. - """ - from pydantic import BaseModel - - class _Usage(BaseModel): - total_tokens: int = 30 - prompt_tokens: int = 20 - completion_tokens: int = 10 - - class _Resp(BaseModel): - usage: _Usage = _Usage() - _hidden_params: dict = {} - - router = Router(model_list=model_list) - router.get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-remaining-tokens": 1000, - "x-ratelimit-remaining-requests": 100, "x-ratelimit-remaining-input-tokens": 1000, "x-ratelimit-remaining-output-tokens": 500, } @@ -1005,14 +960,155 @@ async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_l await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") headers = resp._hidden_params["additional_headers"] - # TPM/RPM headers replay the in-flight increment... - assert headers["x-ratelimit-remaining-tokens"] == 970 + assert headers["x-ratelimit-remaining-tokens"] == 958 assert headers["x-ratelimit-remaining-requests"] == 99 - # ...but the reservation-based input/output headers pass through unchanged. + assert headers["x-ratelimit-limit-tokens"] == 1000 + assert headers["x-ratelimit-limit-requests"] == 100 assert headers["x-ratelimit-remaining-input-tokens"] == 1000 assert headers["x-ratelimit-remaining-output-tokens"] == 500 +def _rpm_tpm_router(model_id: str) -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake", "tpm": 1000, "rpm": 100}, + "model_info": {"id": model_id}, + } + ] + ) + + +def _ratelimit_headers(response) -> dict: + return {k: v for k, v in response._hidden_params["additional_headers"].items() if k.startswith("x-ratelimit-")} + + +@pytest.mark.asyncio +async def test_acompletion_headers_read_post_increment_counter_and_count_once(): + """ + LIT-3058 regression: the remaining-* headers on the response must already include + this request, and the request must land in the router counter exactly once even + though `deployment_callback_on_success` still runs after the response returns. + """ + router = _rpm_tpm_router("lit-3058-async") + + response = await router.acompletion( + model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong" + ) + total_tokens = response.usage.total_tokens + assert total_tokens > 0 + + headers = _ratelimit_headers(response) + assert headers["x-ratelimit-remaining-tokens"] == 1000 - total_tokens + assert headers["x-ratelimit-remaining-requests"] == 99 + assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1) + + await asyncio.sleep(0.5) + assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1) + + +@pytest.mark.asyncio +async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deployment_name(): + """The counter key is written under the resolved model name, which is what the usage reader looks up.""" + router = Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "sk-fake", "tpm": 1000, "rpm": 100}, + "model_info": {"id": "lit-3058-wildcard"}, + } + ] + ) + + response = await router.acompletion( + model="openai/gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong" + ) + total_tokens = response.usage.total_tokens + + headers = _ratelimit_headers(response) + assert headers["x-ratelimit-remaining-tokens"] == 1000 - total_tokens + assert headers["x-ratelimit-remaining-requests"] == 99 + assert await router.get_model_group_usage("openai/gpt-5-mini") == (total_tokens, 1) + + +@pytest.mark.asyncio +async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_on_completion(): + """ + A stream has no usage when the headers are built: the request is counted before the + headers and the final token usage is added once when the stream completes. + """ + router = _rpm_tpm_router("lit-3058-stream") + + stream = await router.acompletion( + model="gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + mock_response="pong pong pong", + stream=True, + stream_options={"include_usage": True}, + ) + headers = _ratelimit_headers(stream) + assert headers["x-ratelimit-remaining-tokens"] == 1000 + assert headers["x-ratelimit-remaining-requests"] == 99 + assert await router.get_model_group_usage("gpt-5-mini") == (None, 1) + + chunks = [chunk async for chunk in stream] + total_tokens = chunks[-1].usage.total_tokens + assert total_tokens > 0 + + await asyncio.sleep(0.5) + assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1) + + +@pytest.mark.asyncio +async def test_deployment_callback_on_success_adds_only_uncounted_tokens(): + """ + When the request was already counted before the headers, the success callback adds + only the tokens that were not known at that point and never a second request. + """ + import time + + router = _rpm_tpm_router("lit-3058-callback") + standard_logging_payload = create_standard_logging_payload() + standard_logging_payload["total_tokens"] = 100 + kwargs = { + "litellm_params": { + "metadata": { + "deployment": "gpt-5-mini", + "model_group": "gpt-5-mini", + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 60, + }, + "model_info": {"id": "lit-3058-callback"}, + }, + "standard_logging_object": standard_logging_payload, + } + + tpm_key = await router.deployment_callback_on_success( + kwargs=kwargs, + completion_response=litellm.ModelResponse(model="gpt-5-mini", usage={"total_tokens": 100}), + start_time=time.time(), + end_time=time.time(), + ) + + assert tpm_key is not None + assert await router.get_model_group_usage("gpt-5-mini") == (40, None) + + +@pytest.mark.asyncio +async def test_increment_deployment_usage_for_response_skips_session_wrappers(): + """WebSocket and realtime session wrappers return None and are not counted as a request.""" + router = _rpm_tpm_router("lit-3058-ws") + request_kwargs = { + "model": "gpt-5-mini", + "litellm_metadata": {"model_group": "gpt-5-mini", "model_info": {"id": "lit-3058-ws"}}, + } + + await router.increment_deployment_usage_for_response(response=None, request_kwargs=request_kwargs) + + assert await router.get_model_group_usage("gpt-5-mini") == (None, None) + assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in request_kwargs["litellm_metadata"] + + @pytest.mark.asyncio async def test_get_model_group_io_token_usage_sums_across_deployments(): """ @@ -1154,8 +1250,8 @@ async def test_set_response_headers_native_input_token_header_does_not_suppress_ await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") headers = resp._hidden_params["additional_headers"] - assert headers["x-ratelimit-remaining-tokens"] == 958 - assert headers["x-ratelimit-remaining-requests"] == 99 + assert headers["x-ratelimit-remaining-tokens"] == 1000 + assert headers["x-ratelimit-remaining-requests"] == 100 # the provider's native header is left untouched assert headers["x-ratelimit-remaining-input-tokens"] == 5 @@ -1187,7 +1283,7 @@ async def test_set_response_headers_native_token_header_does_not_suppress_io_hea headers = resp._hidden_params["additional_headers"] assert headers["x-ratelimit-remaining-tokens"] == 5 - assert headers["x-ratelimit-remaining-requests"] == 99 + assert headers["x-ratelimit-remaining-requests"] == 100 assert headers["x-ratelimit-remaining-input-tokens"] == 900 assert headers["x-ratelimit-remaining-output-tokens"] == 450 @@ -1196,8 +1292,7 @@ async def test_set_response_headers_native_token_header_does_not_suppress_io_hea async def test_set_response_headers_handles_missing_usage(model_list): """ Streaming chunks and some response shapes may lack a `usage` attribute or - populated `total_tokens`. The in-flight subtraction must default to 0 - tokens (still subtract 1 from requests) and never raise. + populated `total_tokens`. Header composition must not depend on usage and never raise. """ from pydantic import BaseModel @@ -1218,7 +1313,7 @@ async def test_set_response_headers_handles_missing_usage(model_list): headers = resp._hidden_params["additional_headers"] assert headers["x-ratelimit-remaining-tokens"] == 1000 - assert headers["x-ratelimit-remaining-requests"] == 99 + assert headers["x-ratelimit-remaining-requests"] == 100 @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 099afa57eec..a7b42841e53 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -43,7 +43,11 @@ from litellm.litellm_core_utils.get_provider_specific_headers import ( from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, ) -from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY +from litellm.constants import ( + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, + SESSION_ID_OMITTED_METADATA_KEY, +) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import CredentialItem @@ -7354,6 +7358,7 @@ _PLANTED_STAMPS = { "original_model_group": "spoofed-group", "request_retry_count": -100, "_client_output_ceiling": {"api_base": "https://attacker.example"}, + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 10**9, "client_key": "client_value", } @@ -7386,6 +7391,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "original_model_group" not in updated["metadata"] assert "_client_output_ceiling" not in updated["metadata"] assert "request_retry_count" not in updated["metadata"] + assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" From 2bb478fb59b99d3ea46e473bad44f8c2157b3430 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 18:45:11 +0000 Subject: [PATCH 101/187] fix(ui): keep http and upper-case https skill sources clickable on the detail page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/components/claude_code_plugins/helpers.test.ts | 9 +++++++++ .../src/components/claude_code_plugins/helpers.ts | 3 ++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts index 5bf782aabf3..e40e0e0c783 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts @@ -157,6 +157,15 @@ describe("getSourceLink", () => { expect(getSourceLink({ source: "github" })).toBeNull(); }); + it("keeps http and upper-case https urls registered through the api clickable", () => { + expect(getSourceLink({ source: "url", url: "http://git.internal.example/org/repo" })).toBe( + "http://git.internal.example/org/repo", + ); + expect(getSourceLink({ source: "git-subdir", url: "HTTPS://gitlab.com/org/repo", path: "sub/dir" })).toBe( + "HTTPS://gitlab.com/org/repo", + ); + }); + it("returns null for an ssh clone url, which is not browsable", () => { expect(getSourceLink({ source: "url", url: "git@ghe.example.com:org/repo.git" })).toBeNull(); expect(getSourceLink({ source: "url", url: "ssh://git@ghe.example.com/org/repo.git" })).toBeNull(); diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts index b4d9eab08ec..8cf620d9077 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts @@ -37,6 +37,7 @@ const IPV4_HOST_REGEX = /^\d{1,3}(\.\d{1,3}){3}$/; const GITHUB_ORG_REGEX = /^[A-Za-z0-9-]+$/; const GITHUB_REPO_REGEX = /^[A-Za-z0-9._-]+$/; +const BROWSABLE_URL_REGEX = /^https?:\/\//i; const SSH_SCHEME = "ssh://"; const SSH_SCP_REGEX = /^([a-z0-9._-]+)@([^:/@]+):(?!\/)(.+)$/i; @@ -316,7 +317,7 @@ export const getSourceLink = (source: PluginSource): string | null => { return `https://github.com/${source.repo}`; } const linksToUrl = source.source === "url" || source.source === "git-subdir" || source.source === "archive"; - return linksToUrl && source.url?.startsWith("https://") ? source.url : null; + return linksToUrl && source.url && BROWSABLE_URL_REGEX.test(source.url) ? source.url : null; }; /** From 99aa9f76c8fc0a84969a975562dd04bbd3a7b160 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 18:51:22 +0000 Subject: [PATCH 102/187] fix(proxy): build failure headers immutably to keep LIT002 within budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/batches_endpoints/endpoints.py | 3 ++- litellm/proxy/common_utils/openai_error_payload.py | 7 +++++++ litellm/proxy/proxy_server.py | 4 ++-- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index f37c06aea97..5d9ecddd4c2 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -7,6 +7,7 @@ import asyncio import os from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final, cast from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response @@ -731,7 +732,7 @@ async def list_batches( ) verbose_proxy_logger.debug("GET /v1/batches after=%s limit=%s", after, limit) - data: dict = {} + data: Mapping[str, object] = MappingProxyType({}) try: if llm_router is None: raise HTTPException( diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index cbc8c78d4f9..202c61b620e 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -68,3 +68,10 @@ def with_litellm_call_id(exc: ProxyException, litellm_call_id: str | None) -> Pr if litellm_call_id is not None: exc.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id) return exc + + +def headers_with_litellm_call_id(headers: Mapping[str, str] | None, litellm_call_id: str) -> Mapping[str, str]: + """``headers`` plus ``x-litellm-call-id``, keeping the value they already carry under that name.""" + if headers is None: + return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id}) + return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id, **headers}) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f72cabf076..45beb8cc93c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -392,7 +392,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) from litellm.proxy.common_utils.openai_error_payload import ( - LITELLM_CALL_ID_HEADER, + headers_with_litellm_call_id, litellm_call_id_headers, with_litellm_call_id, ) @@ -11696,7 +11696,7 @@ async def audio_speech( raise HTTPException( status_code=e.status_code, detail=e.detail, - headers={LITELLM_CALL_ID_HEADER: litellm_call_id, **(e.headers or {})}, + headers=headers_with_litellm_call_id(e.headers, litellm_call_id), ) raise ProxyException( message=getattr(e, "message", f"{e}"), From a2108f02fb83bac626439dfd8f76b41faa57e695 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 18:51:40 +0000 Subject: [PATCH 103/187] test(router): cover _increment_deployment_usage delta and unlimited deployment behavior Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_router_helper_utils.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 215101735c8..a7bdb754c2c 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1109,6 +1109,57 @@ async def test_increment_deployment_usage_for_response_skips_session_wrappers(): assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in request_kwargs["litellm_metadata"] +@pytest.mark.asyncio +async def test_increment_deployment_usage_writes_only_positive_deltas_for_limited_deployments(): + router = _rpm_tpm_router("lit-3058-delta") + unlimited = Router( + model_list=[ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"}, + "model_info": {"id": "lit-3058-unlimited"}, + } + ] + ) + + tpm_key = await router._increment_deployment_usage( + deployment_id="lit-3058-delta", + deployment_name="gpt-5-mini", + model_group="gpt-5-mini", + total_tokens=25, + rpm_increment=1, + parent_otel_span=None, + ) + assert tpm_key is not None + assert await router.get_model_group_usage("gpt-5-mini") == (25, 1) + + assert ( + await router._increment_deployment_usage( + deployment_id="lit-3058-delta", + deployment_name="gpt-5-mini", + model_group="gpt-5-mini", + total_tokens=0, + rpm_increment=0, + parent_otel_span=None, + ) + is None + ) + assert await router.get_model_group_usage("gpt-5-mini") == (25, 1) + + assert ( + await unlimited._increment_deployment_usage( + deployment_id="lit-3058-unlimited", + deployment_name="gpt-5-mini", + model_group="gpt-5-mini", + total_tokens=25, + rpm_increment=1, + parent_otel_span=None, + ) + is None + ) + assert await unlimited.get_model_group_usage("gpt-5-mini") == (None, None) + + @pytest.mark.asyncio async def test_get_model_group_io_token_usage_sums_across_deployments(): """ From ad25fe1886e0444b32d9acef2844c94eb7e4f20d Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 18:56:56 +0000 Subject: [PATCH 104/187] fix(guardrails): hold unscannable Responses windows and key terminal envelopes by output items Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_translation/handler.py | 6 +- .../unified_guardrail/unified_guardrail.py | 4 +- ...test_openai_responses_guardrail_handler.py | 18 +++ .../test_streaming_buffer_until_moderated.py | 125 ++++++++++++++++++ 4 files changed, 147 insertions(+), 6 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 321f0a0d6d3..1ab4811b8a0 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1175,13 +1175,11 @@ class OpenAIResponsesHandler(BaseTranslation): last_event_type: Final = stream_item_field(last_event, "type") if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value: return None - if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value: + if last_event_type in _TERMINAL_ENVELOPE_EVENT_TYPES: return self._completed_response_scan_key(stream_item_field(last_event, "response")) - stream_ended: Final = self._check_streaming_has_ended(responses_so_far) return StreamingScanKey( texts=(self.get_streaming_string_so_far(responses_so_far),), - stream_ended=stream_ended, - tool_calls_in_flight=not stream_ended and self._has_streamed_tool_call_events(responses_so_far), + tool_calls_in_flight=self._has_streamed_tool_call_events(responses_so_far), ) @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 5ac9070faa6..15757681d38 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -1082,7 +1082,7 @@ class UnifiedLLMGuardrails(CustomLogger): scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far) if scan_key is not None: tool_calls_in_flight = scan_key.tool_calls_in_flight - hold_window = buffer_until_moderated and tool_calls_in_flight + hold_window = buffer_until_moderated and (scan_key is None or tool_calls_in_flight) if _is_redundant_scan(scan_key, last_scan_key): verbose_proxy_logger.debug( "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round", @@ -1159,7 +1159,7 @@ class UnifiedLLMGuardrails(CustomLogger): last_scan_key = scan_key if hold_window: verbose_proxy_logger.debug( - "Holding %s buffered chunks for guardrail %s: streamed tool calls await the end-of-stream scan", + "Holding %s buffered chunks for guardrail %s: this round could not scan the whole window", len(withheld_items), guardrail_to_apply.guardrail_name, ) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 60375b03404..c714b5d378a 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -3212,6 +3212,24 @@ class TestOpenAIResponsesHandlerStreamingScanKey: done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}} assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None + @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"]) + def test_non_completed_terminal_envelopes_key_their_output_items(self, terminal_type): + handler = OpenAIResponsesHandler() + arguments_delta = { + "type": "response.function_call_arguments.delta", + "sequence_number": 1, + "item_id": "fc_1", + "delta": '{"city":', + } + function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city":'} + terminal = {"type": terminal_type, "sequence_number": 2, "response": {"id": "resp_1", "output": [function_call]}} + mid_stream_key = handler.get_streaming_scan_key([arguments_delta]) + ended_key = handler.get_streaming_scan_key([arguments_delta, terminal]) + assert ended_key.stream_ended is True + assert ended_key.tool_calls_in_flight is False + assert len(ended_key.tool_calls) == 1 + assert ended_key != mid_stream_key + def test_streamed_tool_call_events_flag_tool_calls_in_flight_until_the_stream_ends(self): handler = OpenAIResponsesHandler() added = { diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py index cbc17e680f0..ed548155b60 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py @@ -121,6 +121,32 @@ class _SecondScanBlockingGuardrail(_CountingPassingGuardrail): return inputs +class _MarkerBlockingGuardrail(_CountingPassingGuardrail): + """Blocks as soon as the inspected input field (texts or tool_calls) carries the marker.""" + + def __init__(self, *args, marker: str, field: Literal["texts", "tool_calls"] = "texts", **kwargs): + super().__init__(*args, **kwargs) + self.marker = marker + self.field = field + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.scan_count += 1 + if self.marker in json.dumps(inputs.get(self.field, [])): + raise ModifyResponseException( + message=BLOCK_MESSAGE, + model="gpt-4o", + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + return inputs + + def _sse_event(event_type: str, data: dict) -> bytes: return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() @@ -271,6 +297,76 @@ async def _run_windowed( return collected, yielded_count +def _responses_message_stream_events(text_chunks: List[str]) -> List[dict]: + message = {"type": "message", "id": "msg_1", "status": "completed", "role": "assistant"} + content = [{"type": "output_text", "text": "".join(text_chunks), "annotations": []}] + return [ + {"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}}, + *( + {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": text} + for text in text_chunks + ), + {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": content}}, + { + "type": "response.completed", + "response": {"id": "resp_1", "model": "gpt-4o", "status": "completed", "output": [{**message, "content": content}]}, + }, + ] + + +def _responses_truncated_function_call_events(text: str, argument_chunks: List[str]) -> List[dict]: + message = {"type": "message", "id": "msg_1", "status": "completed", "role": "assistant"} + content = [{"type": "output_text", "text": text, "annotations": []}] + function_call = {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "run_shell"} + return [ + {"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}}, + {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": text}, + {"type": "response.output_item.added", "output_index": 1, "item": {**function_call, "arguments": ""}}, + *( + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": arguments} + for arguments in argument_chunks + ), + { + "type": "response.incomplete", + "response": { + "id": "resp_1", + "model": "gpt-4o", + "status": "incomplete", + "output": [ + {**message, "content": content}, + {**function_call, "arguments": "".join(argument_chunks), "status": "incomplete"}, + ], + }, + }, + ] + + +async def _replay(events: List[dict]) -> AsyncGenerator[dict, None]: + for event in events: + yield event + + +async def _run_windowed_responses(guardrail: CustomGuardrail, events: List[dict]) -> str: + guardrail.streaming_buffer_until_moderated = True + guardrail.streaming_buffer_release_on_scan = True + guardrail.streaming_sampling_rate = 2 + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses") + request_data = { + "input": "hi", + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: List[Any] = [] + async for chunk in unified.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_replay(events), + request_data=request_data, + ): + collected.append(chunk) + return json.dumps([chunk if isinstance(chunk, dict) else str(chunk) for chunk in collected]) + + def _chat_text(chunks: List[Any]) -> str: return "".join( choice.delta.content or "" @@ -365,6 +461,35 @@ async def test_windowed_buffer_holds_tool_call_windows_until_end_of_stream_scan( assert guardrail.tool_call_scan_indexes == [guardrail.scan_count] +@pytest.mark.asyncio +async def test_windowed_responses_output_item_done_round_keeps_text_window_withheld(): + guardrail = _MarkerBlockingGuardrail(guardrail_name="windowed-responses", event_hook="post_call", marker=ORIGINAL_MARKER) + events = _responses_message_stream_events(["one ", f"{ORIGINAL_MARKER} "]) + + raw = await _run_windowed_responses(guardrail, events) + + assert ORIGINAL_MARKER not in raw, f"unscanned window leaked: {raw!r}" + assert BLOCK_MESSAGE in raw + assert guardrail.scan_count >= 1 + + +@pytest.mark.asyncio +async def test_windowed_responses_incomplete_stream_scans_tool_call_before_release(): + guardrail = _MarkerBlockingGuardrail( + guardrail_name="windowed-responses-tools", + event_hook="post_call", + marker=TOOL_ARGUMENTS_MARKER, + field="tool_calls", + ) + events = _responses_truncated_function_call_events("hi ", ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}']) + + raw = await _run_windowed_responses(guardrail, events) + + assert '"hi "' in raw + assert TOOL_ARGUMENTS_MARKER not in raw, f"unscanned tool call leaked: {raw!r}" + assert BLOCK_MESSAGE in raw + + @pytest.mark.asyncio async def test_windowed_buffer_with_explicit_end_of_stream_only_stays_fully_buffered(): guardrail = _CountingPassingGuardrail(guardrail_name="windowed-eos", event_hook="post_call") From e105cde56a585ed4c49c3013a8da2c3e3b8f6066 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 19:02:15 +0000 Subject: [PATCH 105/187] fix(router): sync worker-local usage cache to shared post-increment counters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 7 +++- .../test_router_helper_utils.py | 40 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index f8abd9be758..71588109c45 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8284,10 +8284,15 @@ class Router: for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment)) if increment_value > 0 ] - await self.cache.async_increment_cache_pipeline( + post_increment_values: Final = await self.cache.async_increment_cache_pipeline( increment_list=pipeline_operations, parent_otel_span=parent_otel_span, ) + if post_increment_values is not None and self.cache.redis_cache is not None: + for operation, value in zip(pipeline_operations, post_increment_values): + await self.cache.async_set_cache( + operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value + ) return tpm_key def sync_deployment_callback_on_success( diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index a7bdb754c2c..8646843c0de 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1160,6 +1160,46 @@ async def test_increment_deployment_usage_writes_only_positive_deltas_for_limite assert await unlimited.get_model_group_usage("gpt-5-mini") == (None, None) +def _shared_redis_stub(store: dict) -> MagicMock: + from litellm.caching.redis_cache import RedisCache + + async def increment_pipeline(increment_list, **kwargs): + for op in increment_list: + store[op["key"]] = store.get(op["key"], 0.0) + op["increment_value"] + return [store[op["key"]] for op in increment_list] + + async def batch_get(keys, **kwargs): + return {key: store.get(key) for key in keys} + + redis_stub = MagicMock(spec=RedisCache) + redis_stub.async_increment_pipeline = increment_pipeline + redis_stub.async_batch_get_cache = batch_get + return redis_stub + + +@pytest.mark.asyncio +async def test_headers_on_fresh_worker_reflect_shared_redis_usage(): + from litellm.caching.dual_cache import DualCache + from litellm.caching.in_memory_cache import InMemoryCache + + store: dict = {} + worker_a = _rpm_tpm_router("lit-3058-workers") + worker_b = _rpm_tpm_router("lit-3058-workers") + worker_a.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache()) + worker_b.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache()) + + messages = [{"role": "user", "content": "hi"}] + tokens_on_a = 0 + for _ in range(3): + response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") + tokens_on_a += response.usage.total_tokens + + response = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") + headers = _ratelimit_headers(response) + assert headers["x-ratelimit-remaining-requests"] == 96 + assert headers["x-ratelimit-remaining-tokens"] == 1000 - tokens_on_a - response.usage.total_tokens + + @pytest.mark.asyncio async def test_get_model_group_io_token_usage_sums_across_deployments(): """ From 7734e3e186cb9194db22c6a4975c664a40b5a685 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 19:17:13 +0000 Subject: [PATCH 106/187] fix(models): sonnet 4.5 1M input, daybreak alias, Mistral GLM 5.3, Azure dated snapshots, Together/OpenRouter sync Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 337 +++++++++++++++--- model_prices_and_context_window.json | 337 +++++++++++++++--- 2 files changed, 568 insertions(+), 106 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7feef3be8ea..231169b4c9f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -7566,6 +7566,70 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-sol-2026-07-09": { + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_priority": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_flex": 2.5e-06, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_priority": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, + "cache_read_input_token_cost_flex": 2e-07, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, + "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_priority": 8e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, + "input_cost_per_token_flex": 2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_priority": 4e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, + "output_cost_per_token_flex": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure/gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -7630,6 +7694,70 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-terra-2026-07-09": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_priority": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, + "cache_creation_input_token_cost_flex": 1.25e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_read_input_token_cost_priority": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, + "cache_read_input_token_cost_flex": 1e-07, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "input_cost_per_token_priority": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 8e-06, + "input_cost_per_token_flex": 1e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "output_cost_per_token_above_272k_tokens_flex": 9e-06, + "output_cost_per_token_priority": 2.4e-05, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, + "output_cost_per_token_flex": 6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure/gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, @@ -7798,6 +7926,54 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-6-astra-2026-09-03": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/gpt-chat-latest": { "cache_read_input_token_cost": 5e-07, "deprecation_date": "2026-12-02", @@ -14186,7 +14362,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", - "max_input_tokens": 200000, + "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -14221,7 +14397,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", - "max_input_tokens": 200000, + "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -14242,7 +14418,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.anthropic.com/en/docs/about-claude/pricing" }, "claude-sonnet-5": { "deprecation_date": "2027-06-30", @@ -32459,7 +32636,7 @@ "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "openai", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -32498,7 +32675,7 @@ "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "openai", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -36833,6 +37010,57 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/zai-glm-5-3": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/zai-glm-5": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/zai-glm-latest": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/glm-5-2": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, @@ -40578,15 +40806,16 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-chat": { - "input_cost_per_token": 3.2e-07, + "input_cost_per_token": 2.574e-07, "litellm_provider": "openrouter", "max_input_tokens": 65536, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 8.9e-07, + "output_cost_per_token": 1.0287e-06, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/deepseek/deepseek-chat-v3-0324": { "input_cost_per_token": 2.5e-07, @@ -40679,21 +40908,21 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.59908e-07, + "input_cost_per_token": 1.6e-06, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.719816e-06, + "output_cost_per_token": 3.2e-06, "source": "https://openrouter.ai/deepseek/deepseek-v4-pro", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.1659e-08 + "cache_read_input_token_cost": 1.35e-07 }, "openrouter/deepseek/deepseek-v4.1-flash": { "input_cost_per_token": 1.5e-07, @@ -40988,12 +41217,13 @@ "supports_vision": true }, "openrouter/gryphe/mythomax-l2-13b": { - "input_cost_per_token": 6e-08, + "input_cost_per_token": 8e-08, "litellm_provider": "openrouter", "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 6e-08, - "supports_tool_choice": true + "output_cost_per_token": 1.1e-07, + "supports_tool_choice": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/mancer/weaver": { "input_cost_per_token": 4e-07, @@ -41129,14 +41359,15 @@ "max_output_tokens": 131072 }, "openrouter/mistralai/mistral-small-3.2-24b-instruct": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 9.375e-08, "litellm_provider": "openrouter", "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2e-07, + "output_cost_per_token": 2.5e-07, "supports_tool_choice": true, "max_input_tokens": 128000, - "max_output_tokens": 128000 + "max_output_tokens": 128000, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/mistralai/mixtral-8x22b-instruct": { "input_cost_per_token": 2e-06, @@ -41657,13 +41888,13 @@ "supports_tool_choice": true }, "openrouter/qwen/qwen3-235b-a22b-2507": { - "input_cost_per_token": 2.2e-07, + "input_cost_per_token": 8.75e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 8.8e-07, + "output_cost_per_token": 3.5e-07, "source": "https://openrouter.ai/qwen/qwen3-235b-a22b-2507", "supports_function_calling": true, "supports_tool_choice": true @@ -41696,13 +41927,13 @@ "supports_vision": true }, "openrouter/qwen/qwen3.5-35b-a3b": { - "input_cost_per_token": 3.125e-07, + "input_cost_per_token": 1.625e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1.25e-06, + "output_cost_per_token": 1.3e-06, "source": "https://openrouter.ai/qwen/qwen3.5-35b-a3b", "supports_function_calling": true, "supports_reasoning": true, @@ -44718,7 +44949,7 @@ "deprecation_date": "2026-06-04", "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 2e-06, "source": "https://api.together.ai/v1/models", @@ -44788,7 +45019,7 @@ "supports_parallel_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 16384 }, "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": { @@ -44971,7 +45202,7 @@ "deprecation_date": "2026-04-02", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.1e-06, "source": "https://api.together.ai/v1/models", @@ -64723,9 +64954,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.3-flash": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 5e-07, - "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3e-07, + "cache_read_input_token_cost": 1.8e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -64775,9 +65006,9 @@ "supports_prompt_caching": true }, "openrouter/qwen/qwen3.8-27b": { - "input_cost_per_token": 4.2e-07, - "output_cost_per_token": 3e-06, - "cache_read_input_token_cost": 8.5e-08, + "input_cost_per_token": 2.14e-07, + "output_cost_per_token": 2.55e-06, + "cache_read_input_token_cost": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 131072, @@ -64860,9 +65091,9 @@ "supports_prompt_caching": true }, "openrouter/deepseek/deepseek-v4-flash-0731": { - "input_cost_per_token": 6.5e-08, - "output_cost_per_token": 1.8e-07, - "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 1.2e-07, + "cache_read_input_token_cost": 1.2e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, @@ -64930,9 +65161,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 2.1e-06, - "output_cost_per_token": 1.053e-05, - "cache_read_input_token_cost": 2.35e-07, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -65029,9 +65260,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.2": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2e-06, - "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 1.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 131072, @@ -65062,9 +65293,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k2.7-code": { - "input_cost_per_token": 7.1e-07, - "output_cost_per_token": 3.5e-06, - "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 7.062e-07, + "output_cost_per_token": 3.21e-06, + "cache_read_input_token_cost": 1.8e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 235929, @@ -65366,8 +65597,8 @@ "supports_prompt_caching": true }, "openrouter/google/gemma-4-26b-a4b-it": { - "input_cost_per_token": 4.2e-08, - "output_cost_per_token": 2.2e-07, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -65914,8 +66145,8 @@ "supports_vision": true }, "openrouter/qwen/qwen3-vl-30b-a3b-instruct": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 6e-07, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 5.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -66175,8 +66406,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-30b-a3b-instruct-2507": { - "input_cost_per_token": 9e-08, - "output_cost_per_token": 3e-07, + "input_cost_per_token": 4.815e-08, + "output_cost_per_token": 1.9305e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32000, @@ -66235,7 +66466,7 @@ "supports_vision": false }, "openrouter/minimax/minimax-m1": { - "input_cost_per_token": 5.5e-07, + "input_cost_per_token": 4e-07, "output_cost_per_token": 2.2e-06, "litellm_provider": "openrouter", "max_input_tokens": 1000000, @@ -66374,8 +66605,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-14b": { - "input_cost_per_token": 2.275e-07, - "output_cost_per_token": 9.1e-07, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 2.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 16384, @@ -66437,8 +66668,8 @@ "supports_prompt_caching": true }, "openrouter/meta-llama/llama-4-maverick": { - "input_cost_per_token": 2e-07, - "output_cost_per_token": 6.96e-07, + "input_cost_per_token": 1.875e-07, + "output_cost_per_token": 6.525e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 115200, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7feef3be8ea..231169b4c9f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7566,6 +7566,70 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-sol-2026-07-09": { + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_priority": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_flex": 2.5e-06, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_priority": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, + "cache_read_input_token_cost_flex": 2e-07, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, + "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_priority": 8e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, + "input_cost_per_token_flex": 2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_priority": 4e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, + "output_cost_per_token_flex": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure/gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -7630,6 +7694,70 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-terra-2026-07-09": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_priority": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, + "cache_creation_input_token_cost_flex": 1.25e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_read_input_token_cost_priority": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, + "cache_read_input_token_cost_flex": 1e-07, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "input_cost_per_token_priority": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 8e-06, + "input_cost_per_token_flex": 1e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "output_cost_per_token_above_272k_tokens_flex": 9e-06, + "output_cost_per_token_priority": 2.4e-05, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, + "output_cost_per_token_flex": 6e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure/gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, @@ -7798,6 +7926,54 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-6-astra-2026-09-03": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/gpt-chat-latest": { "cache_read_input_token_cost": 5e-07, "deprecation_date": "2026-12-02", @@ -14186,7 +14362,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", - "max_input_tokens": 200000, + "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -14221,7 +14397,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", - "max_input_tokens": 200000, + "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -14242,7 +14418,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.anthropic.com/en/docs/about-claude/pricing" }, "claude-sonnet-5": { "deprecation_date": "2027-06-30", @@ -32459,7 +32636,7 @@ "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "openai", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -32498,7 +32675,7 @@ "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "openai", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -36833,6 +37010,57 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/zai-glm-5-3": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/zai-glm-5": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/zai-glm-latest": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/glm-5-2": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, @@ -40578,15 +40806,16 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-chat": { - "input_cost_per_token": 3.2e-07, + "input_cost_per_token": 2.574e-07, "litellm_provider": "openrouter", "max_input_tokens": 65536, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 8.9e-07, + "output_cost_per_token": 1.0287e-06, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/deepseek/deepseek-chat-v3-0324": { "input_cost_per_token": 2.5e-07, @@ -40679,21 +40908,21 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.59908e-07, + "input_cost_per_token": 1.6e-06, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.719816e-06, + "output_cost_per_token": 3.2e-06, "source": "https://openrouter.ai/deepseek/deepseek-v4-pro", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.1659e-08 + "cache_read_input_token_cost": 1.35e-07 }, "openrouter/deepseek/deepseek-v4.1-flash": { "input_cost_per_token": 1.5e-07, @@ -40988,12 +41217,13 @@ "supports_vision": true }, "openrouter/gryphe/mythomax-l2-13b": { - "input_cost_per_token": 6e-08, + "input_cost_per_token": 8e-08, "litellm_provider": "openrouter", "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 6e-08, - "supports_tool_choice": true + "output_cost_per_token": 1.1e-07, + "supports_tool_choice": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/mancer/weaver": { "input_cost_per_token": 4e-07, @@ -41129,14 +41359,15 @@ "max_output_tokens": 131072 }, "openrouter/mistralai/mistral-small-3.2-24b-instruct": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 9.375e-08, "litellm_provider": "openrouter", "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2e-07, + "output_cost_per_token": 2.5e-07, "supports_tool_choice": true, "max_input_tokens": 128000, - "max_output_tokens": 128000 + "max_output_tokens": 128000, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/mistralai/mixtral-8x22b-instruct": { "input_cost_per_token": 2e-06, @@ -41657,13 +41888,13 @@ "supports_tool_choice": true }, "openrouter/qwen/qwen3-235b-a22b-2507": { - "input_cost_per_token": 2.2e-07, + "input_cost_per_token": 8.75e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 8.8e-07, + "output_cost_per_token": 3.5e-07, "source": "https://openrouter.ai/qwen/qwen3-235b-a22b-2507", "supports_function_calling": true, "supports_tool_choice": true @@ -41696,13 +41927,13 @@ "supports_vision": true }, "openrouter/qwen/qwen3.5-35b-a3b": { - "input_cost_per_token": 3.125e-07, + "input_cost_per_token": 1.625e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1.25e-06, + "output_cost_per_token": 1.3e-06, "source": "https://openrouter.ai/qwen/qwen3.5-35b-a3b", "supports_function_calling": true, "supports_reasoning": true, @@ -44718,7 +44949,7 @@ "deprecation_date": "2026-06-04", "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 2e-06, "source": "https://api.together.ai/v1/models", @@ -44788,7 +45019,7 @@ "supports_parallel_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 16384 }, "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": { @@ -44971,7 +45202,7 @@ "deprecation_date": "2026-04-02", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.1e-06, "source": "https://api.together.ai/v1/models", @@ -64723,9 +64954,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.3-flash": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 5e-07, - "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3e-07, + "cache_read_input_token_cost": 1.8e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -64775,9 +65006,9 @@ "supports_prompt_caching": true }, "openrouter/qwen/qwen3.8-27b": { - "input_cost_per_token": 4.2e-07, - "output_cost_per_token": 3e-06, - "cache_read_input_token_cost": 8.5e-08, + "input_cost_per_token": 2.14e-07, + "output_cost_per_token": 2.55e-06, + "cache_read_input_token_cost": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 131072, @@ -64860,9 +65091,9 @@ "supports_prompt_caching": true }, "openrouter/deepseek/deepseek-v4-flash-0731": { - "input_cost_per_token": 6.5e-08, - "output_cost_per_token": 1.8e-07, - "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 1.2e-07, + "cache_read_input_token_cost": 1.2e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, @@ -64930,9 +65161,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 2.1e-06, - "output_cost_per_token": 1.053e-05, - "cache_read_input_token_cost": 2.35e-07, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -65029,9 +65260,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.2": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2e-06, - "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 1.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 131072, @@ -65062,9 +65293,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k2.7-code": { - "input_cost_per_token": 7.1e-07, - "output_cost_per_token": 3.5e-06, - "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 7.062e-07, + "output_cost_per_token": 3.21e-06, + "cache_read_input_token_cost": 1.8e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 235929, @@ -65366,8 +65597,8 @@ "supports_prompt_caching": true }, "openrouter/google/gemma-4-26b-a4b-it": { - "input_cost_per_token": 4.2e-08, - "output_cost_per_token": 2.2e-07, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -65914,8 +66145,8 @@ "supports_vision": true }, "openrouter/qwen/qwen3-vl-30b-a3b-instruct": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 6e-07, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 5.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -66175,8 +66406,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-30b-a3b-instruct-2507": { - "input_cost_per_token": 9e-08, - "output_cost_per_token": 3e-07, + "input_cost_per_token": 4.815e-08, + "output_cost_per_token": 1.9305e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32000, @@ -66235,7 +66466,7 @@ "supports_vision": false }, "openrouter/minimax/minimax-m1": { - "input_cost_per_token": 5.5e-07, + "input_cost_per_token": 4e-07, "output_cost_per_token": 2.2e-06, "litellm_provider": "openrouter", "max_input_tokens": 1000000, @@ -66374,8 +66605,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-14b": { - "input_cost_per_token": 2.275e-07, - "output_cost_per_token": 9.1e-07, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 2.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 16384, @@ -66437,8 +66668,8 @@ "supports_prompt_caching": true }, "openrouter/meta-llama/llama-4-maverick": { - "input_cost_per_token": 2e-07, - "output_cost_per_token": 6.96e-07, + "input_cost_per_token": 1.875e-07, + "output_cost_per_token": 6.525e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 115200, From 8caf2cb61f33dcb55d5b48d3cb698618bba1526d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 19:17:45 +0000 Subject: [PATCH 107/187] fix(models): flag prompt caching on Vertex and Azure AI Grok rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 8 +++ model_prices_and_context_window.json | 8 +++ ...tex_ai_xai_grok_prompt_caching_metadata.py | 58 +++++++++++++++++++ 3 files changed, 74 insertions(+) create mode 100644 tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 231169b4c9f..dbd42728888 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -11833,6 +11833,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -11851,6 +11852,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -50027,6 +50029,7 @@ "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -50043,6 +50046,7 @@ "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -50063,6 +50067,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -50082,6 +50087,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -50102,6 +50108,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -50121,6 +50128,7 @@ "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 231169b4c9f..dbd42728888 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -11833,6 +11833,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -11851,6 +11852,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -50027,6 +50029,7 @@ "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -50043,6 +50046,7 @@ "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -50063,6 +50067,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -50082,6 +50087,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -50102,6 +50108,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -50121,6 +50128,7 @@ "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, diff --git a/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py b/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py new file mode 100644 index 00000000000..b8b3897f004 --- /dev/null +++ b/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py @@ -0,0 +1,58 @@ +from pathlib import Path +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm import get_model_info +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.utils import supports_prompt_caching + +REPO_ROOT: Final = Path(__file__).parents[2] +MODEL: Final = "vertex_ai/xai/grok-4.6" +GROK_KEY_PREFIXES: Final = ("vertex_ai/xai/grok-", "azure_ai/grok-", "xai/grok-") +COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) +MAIN_COST_MAP: Final = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_COST_MAP: Final = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + + +def _cost_map(path: Path) -> dict[str, dict[str, object]]: + return COST_MAP_ADAPTER.validate_json(path.read_bytes()) + + +def test_grok_entries_with_cache_read_price_advertise_prompt_caching() -> None: + cost_map = _cost_map(MAIN_COST_MAP) + cached_grok_entries = tuple( + (key, entry) + for key, entry in cost_map.items() + if key.startswith(GROK_KEY_PREFIXES) and entry.get("cache_read_input_token_cost") + ) + assert cached_grok_entries, "expected at least one grok entry with a cache read price" + + missing_flag = tuple(key for key, entry in cached_grok_entries if entry.get("supports_prompt_caching") is not True) + assert missing_flag == (), ( + f"grok entries with cache_read_input_token_cost missing supports_prompt_caching: {missing_flag}" + ) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_vertex_ai_grok_4_6_supports_prompt_caching_via_get_model_info() -> None: + routed_model, provider, _, _ = get_llm_provider(model=MODEL) + assert (routed_model, provider) == ("xai/grok-4.6", "vertex_ai") + + info = get_model_info(model=routed_model, custom_llm_provider=provider) + assert info["litellm_provider"] == "vertex_ai" + assert info.get("supports_prompt_caching") is True + + assert supports_prompt_caching(model=MODEL) is True + + +def test_vertex_ai_grok_entries_source_and_backup_match() -> None: + main_map = _cost_map(MAIN_COST_MAP) + backup_map = _cost_map(BACKUP_COST_MAP) + + vertex_grok_keys = tuple(key for key in main_map if key.startswith("vertex_ai/xai/grok-")) + assert vertex_grok_keys, "expected at least one vertex_ai/xai/grok- entry" + + mismatched = tuple(key for key in vertex_grok_keys if backup_map.get(key) != main_map[key]) + assert mismatched == (), f"vertex grok entries differ between source and backup: {mismatched}" From 93eefc5922493d59b8ed08072ff072bc0acf32e2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 19:20:20 +0000 Subject: [PATCH 108/187] fix(models): dedupe Fireworks deprecation keys and sync Azure dated snapshot service tiers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 33 +++++++++++-------- model_prices_and_context_window.json | 33 +++++++++++-------- 2 files changed, 40 insertions(+), 26 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dbd42728888..2575b227bfc 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -7825,17 +7825,23 @@ "azure/gpt-5.6-luna-2026-07-09": { "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_priority": 5e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_flex": 1.25e-07, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-08, "cache_read_input_token_cost_priority": 4e-08, "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, + "cache_read_input_token_cost_flex": 1e-08, "deprecation_date": "2028-01-11", "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, + "input_cost_per_token_above_272k_tokens_flex": 2e-07, "input_cost_per_token_priority": 4e-07, "input_cost_per_token_above_272k_tokens_priority": 8e-07, + "input_cost_per_token_flex": 1e-07, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, @@ -7843,8 +7849,10 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "output_cost_per_token_above_272k_tokens": 1.8e-06, + "output_cost_per_token_above_272k_tokens_flex": 9e-07, "output_cost_per_token_priority": 2.4e-06, "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, + "output_cost_per_token_flex": 6e-07, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -8794,12 +8802,15 @@ "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, - "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_priority": 1.25e-06, "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "cache_read_input_token_cost_flex": 2.5e-07, "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, - "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "input_cost_per_token_priority": 1.25e-05, "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "input_cost_per_token_flex": 2.5e-06, "litellm_provider": "azure", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -8807,7 +8818,7 @@ "mode": "chat", "output_cost_per_token": 3e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, - "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, @@ -8894,7 +8905,7 @@ "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, - "cache_read_input_token_cost_priority": 1.38e-06, + "cache_read_input_token_cost_priority": 1.375e-06, "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, @@ -8991,7 +9002,7 @@ "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, - "cache_read_input_token_cost_priority": 1.38e-06, + "cache_read_input_token_cost_priority": 1.375e-06, "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, @@ -23348,8 +23359,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false, - "deprecation_date": "2026-08-27" + "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, @@ -23737,8 +23747,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false, - "deprecation_date": "2026-08-27" + "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/minimax-m3": { "cache_read_input_token_cost": 6e-08, @@ -23845,8 +23854,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false, - "deprecation_date": "2026-08-27" + "supports_vision": false }, "fireworks_ai/glm-4p7": { "cache_read_input_token_cost": 3e-07, @@ -24065,8 +24073,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false, - "deprecation_date": "2026-08-27" + "supports_vision": false }, "fireworks_ai/minimax-m3": { "cache_read_input_token_cost": 6e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dbd42728888..2575b227bfc 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7825,17 +7825,23 @@ "azure/gpt-5.6-luna-2026-07-09": { "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_priority": 5e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_flex": 1.25e-07, "cache_read_input_token_cost": 2e-08, "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-08, "cache_read_input_token_cost_priority": 4e-08, "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, + "cache_read_input_token_cost_flex": 1e-08, "deprecation_date": "2028-01-11", "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, + "input_cost_per_token_above_272k_tokens_flex": 2e-07, "input_cost_per_token_priority": 4e-07, "input_cost_per_token_above_272k_tokens_priority": 8e-07, + "input_cost_per_token_flex": 1e-07, "litellm_provider": "azure", "max_input_tokens": 922000, "max_output_tokens": 128000, @@ -7843,8 +7849,10 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "output_cost_per_token_above_272k_tokens": 1.8e-06, + "output_cost_per_token_above_272k_tokens_flex": 9e-07, "output_cost_per_token_priority": 2.4e-06, "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, + "output_cost_per_token_flex": 6e-07, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -8794,12 +8802,15 @@ "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_above_272k_tokens": 1e-06, - "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_priority": 1.25e-06, "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "cache_read_input_token_cost_flex": 2.5e-07, "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, - "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "input_cost_per_token_priority": 1.25e-05, "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "input_cost_per_token_flex": 2.5e-06, "litellm_provider": "azure", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -8807,7 +8818,7 @@ "mode": "chat", "output_cost_per_token": 3e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, - "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, @@ -8894,7 +8905,7 @@ "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, - "cache_read_input_token_cost_priority": 1.38e-06, + "cache_read_input_token_cost_priority": 1.375e-06, "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, @@ -8991,7 +9002,7 @@ "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, - "cache_read_input_token_cost_priority": 1.38e-06, + "cache_read_input_token_cost_priority": 1.375e-06, "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, @@ -23348,8 +23359,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false, - "deprecation_date": "2026-08-27" + "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, @@ -23737,8 +23747,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false, - "deprecation_date": "2026-08-27" + "supports_vision": false }, "fireworks_ai/accounts/fireworks/models/minimax-m3": { "cache_read_input_token_cost": 6e-08, @@ -23845,8 +23854,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false, - "deprecation_date": "2026-08-27" + "supports_vision": false }, "fireworks_ai/glm-4p7": { "cache_read_input_token_cost": 3e-07, @@ -24065,8 +24073,7 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": false, - "deprecation_date": "2026-08-27" + "supports_vision": false }, "fireworks_ai/minimax-m3": { "cache_read_input_token_cost": 6e-08, From d5837ab97a91e59e58f24c39979d0af0a827fe0f Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 19:24:14 +0000 Subject: [PATCH 109/187] test(guardrails): pin tool-call-only scan keys as non-empty Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_streaming_buffer_until_moderated.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py index ed548155b60..318488c6d7b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py @@ -19,9 +19,11 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, ) +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, + _is_redundant_scan, ) from litellm.types.utils import ( ChatCompletionDeltaToolCall, @@ -461,6 +463,11 @@ async def test_windowed_buffer_holds_tool_call_windows_until_end_of_stream_scan( assert guardrail.tool_call_scan_indexes == [guardrail.scan_count] +def test_tool_call_only_scan_key_is_not_skipped_as_empty(): + assert _is_redundant_scan(StreamingScanKey(texts=("",)), None) is True + assert _is_redundant_scan(StreamingScanKey(texts=("",), tool_calls=("run_shell:{}",)), None) is False + + @pytest.mark.asyncio async def test_windowed_responses_output_item_done_round_keeps_text_window_withheld(): guardrail = _MarkerBlockingGuardrail(guardrail_name="windowed-responses", event_hook="post_call", marker=ORIGINAL_MARKER) From cb264540343fd04a9abcf8317e04fe693e9acb38 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 19:28:50 +0000 Subject: [PATCH 110/187] test(router): cover success callback racing the pre-header count and drop redundant docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 1 - .../test_router_helper_utils.py | 77 +++++++++++++------ 2 files changed, 53 insertions(+), 25 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 71588109c45..cc11a21d966 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8202,7 +8202,6 @@ class Router: response: object, request_kwargs: dict[str, object], ) -> None: - """Count the request before the headers are read; the success callback adds only what is still missing""" if response is None: return try: diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 8646843c0de..664f15889d3 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -11,7 +11,12 @@ import pytest import litellm from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import ModelResponse, StandardLoggingPayload +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY @@ -930,10 +935,6 @@ async def test_set_response_headers(model_list): @pytest.mark.asyncio async def test_set_response_headers_passes_through_post_increment_counters(model_list): - """ - LIT-3058: `make_call` increments the router's TPM/RPM counter before the headers - are built, so `set_response_headers` writes the remaining values it reads as-is. - """ from pydantic import BaseModel class _Usage(BaseModel): @@ -980,17 +981,12 @@ def _rpm_tpm_router(model_id: str) -> Router: ) -def _ratelimit_headers(response) -> dict: +def _ratelimit_headers(response: ModelResponse | CustomStreamWrapper) -> dict[str, int]: return {k: v for k, v in response._hidden_params["additional_headers"].items() if k.startswith("x-ratelimit-")} @pytest.mark.asyncio async def test_acompletion_headers_read_post_increment_counter_and_count_once(): - """ - LIT-3058 regression: the remaining-* headers on the response must already include - this request, and the request must land in the router counter exactly once even - though `deployment_callback_on_success` still runs after the response returns. - """ router = _rpm_tpm_router("lit-3058-async") response = await router.acompletion( @@ -1010,7 +1006,6 @@ async def test_acompletion_headers_read_post_increment_counter_and_count_once(): @pytest.mark.asyncio async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deployment_name(): - """The counter key is written under the resolved model name, which is what the usage reader looks up.""" router = Router( model_list=[ { @@ -1034,10 +1029,6 @@ async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deplo @pytest.mark.asyncio async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_on_completion(): - """ - A stream has no usage when the headers are built: the request is counted before the - headers and the final token usage is added once when the stream completes. - """ router = _rpm_tpm_router("lit-3058-stream") stream = await router.acompletion( @@ -1062,10 +1053,6 @@ async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_ @pytest.mark.asyncio async def test_deployment_callback_on_success_adds_only_uncounted_tokens(): - """ - When the request was already counted before the headers, the success callback adds - only the tokens that were not known at that point and never a second request. - """ import time router = _rpm_tpm_router("lit-3058-callback") @@ -1094,9 +1081,54 @@ async def test_deployment_callback_on_success_adds_only_uncounted_tokens(): assert await router.get_model_group_usage("gpt-5-mini") == (40, None) +class _GatedIncrementCache(DualCache): + def __init__(self) -> None: + super().__init__(in_memory_cache=InMemoryCache()) + self.first_increment_started = asyncio.Event() + self.release_first_increment = asyncio.Event() + self.increment_calls = 0 + + async def async_increment_cache_pipeline( + self, + increment_list: list[RedisPipelineIncrementOperation], + local_only: bool = False, + parent_otel_span: object = None, + **kwargs: object, + ) -> list[float] | None: + self.increment_calls += 1 + if self.increment_calls == 1: + self.first_increment_started.set() + await self.release_first_increment.wait() + return await super().async_increment_cache_pipeline( + increment_list, local_only=local_only, parent_otel_span=parent_otel_span, **kwargs + ) + + +@pytest.mark.asyncio +async def test_success_callback_running_during_pre_header_increment_does_not_double_count(): + router = _rpm_tpm_router("lit-3058-race") + cache = _GatedIncrementCache() + router.cache = cache + + request = asyncio.ensure_future( + router.acompletion(model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong") + ) + await asyncio.wait_for(cache.first_increment_started.wait(), timeout=5) + for _ in range(50): + if get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1: + break + await asyncio.sleep(0.1) + assert get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1 + assert cache.increment_calls == 1 + + cache.release_first_increment.set() + response = await request + + assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1) + + @pytest.mark.asyncio async def test_increment_deployment_usage_for_response_skips_session_wrappers(): - """WebSocket and realtime session wrappers return None and are not counted as a request.""" router = _rpm_tpm_router("lit-3058-ws") request_kwargs = { "model": "gpt-5-mini", @@ -1179,9 +1211,6 @@ def _shared_redis_stub(store: dict) -> MagicMock: @pytest.mark.asyncio async def test_headers_on_fresh_worker_reflect_shared_redis_usage(): - from litellm.caching.dual_cache import DualCache - from litellm.caching.in_memory_cache import InMemoryCache - store: dict = {} worker_a = _rpm_tpm_router("lit-3058-workers") worker_b = _rpm_tpm_router("lit-3058-workers") From 1e169bb912c9a28b3577aca0e0115caa686299b7 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 19:33:07 +0000 Subject: [PATCH 111/187] refactor(guardrails): mark release_on_scan flag as Final Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/unified_guardrail/unified_guardrail.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 15757681d38..7f51d733d4c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -956,7 +956,7 @@ class UnifiedLLMGuardrails(CustomLogger): buffer_until_moderated: bool = _streaming_flag( "streaming_buffer_until_moderated", buffer_until_moderated_default ) - release_on_scan: bool = _streaming_flag("streaming_buffer_release_on_scan", False) + release_on_scan: Final[bool] = _streaming_flag("streaming_buffer_release_on_scan", False) if ( buffer_until_moderated From 39acea0754e0dd5f291e8f99126b1e0b3f5505b3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 12:42:15 -0700 Subject: [PATCH 112/187] feat(e2e): make Claude Code send the same bytes every build The compat cells drove the CLI with a fresh HOME per invocation and the pytest process's own working directory, and both reach the request body. The system prompt names a memory directory built from $CLAUDE_CONFIG_DIR/projects/, so a per-invocation config directory rewrote every body, and the CLI adds a git block for its working directory, so inheriting the checkout rewrote every body once per candidate. The device id churned for the same reason: the CLI mints it once and persists it in .claude.json, which we threw away each call. Nothing here was load-bearing. All three ride in metadata.user_id, whose job is abuse detection, not quota, caching or continuity. So pin the config directory and the working directory at fixed paths, seed the device id, and pin the session id. HOME stays fresh and empty per invocation, so the isolation is no weaker than before, and the CLI's own state no longer outlives the pod either. The working directory is deliberately not the checkout, so a model-directed Read now sees an empty directory rather than the repository. A pinned session id needs --no-session-persistence beside it: the CLI refuses a session id another live process holds, and the matrix runs its cells across xdist workers. Without the flag, six of eight concurrent invocations die on "Session ID is already in use". --- .../test_request_determinism.py | 143 ++++++++++++++++++ tests/e2e/claude_code/cli_driver.py | 57 +++++++ 2 files changed, 200 insertions(+) create mode 100644 tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py new file mode 100644 index 00000000000..09a181162d2 --- /dev/null +++ b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py @@ -0,0 +1,143 @@ +"""The CLI must send the same request bytes from one build to the next. + +Markerless harness test: it drives the real `claude` binary against a local +stub instead of a proxy, so it carries no `e2e` marker. The binary is a +prerequisite of this whole suite, so a missing one is a failure rather than a +skip. + +Two builds differ in ways the driver does not control: a fresh pod, so no CLI +state survives, and a different candidate checked out at a different commit. +Both used to reach the request body, through the memory path the system prompt +names and through the git block the CLI adds for its working directory, so the +shared provider cache missed on every Claude Code cell. This replays those two +differences across a pair of invocations and holds the bytes equal. + +A pinned session id is what makes the second test necessary. The matrix runs +its cells across xdist workers, and the CLI refuses to start a session id that +another live process already holds, so pinning one without also opting out of +session persistence turns most of a parallel run red. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import threading +from collections import Counter +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import List, Tuple + +import pytest + +from claude_code.cli_driver import _stable_cli_state, run_claude +from claude_code.rate_limiter import RateLimiter + +_STUB_REPLY = { + "id": "msg_stub", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 2}, +} + + +def _make_repo(root: Path, subject: str) -> Path: + root.mkdir(parents=True, exist_ok=True) + identity = {"NAME": "t", "EMAIL": "t@e2e"} + env = dict( + os.environ, + **{f"GIT_{role}_{key}": value for role in ("AUTHOR", "COMMITTER") for key, value in identity.items()}, + ) + (root / "file.txt").write_text(subject, encoding="utf-8") + for args in (["init", "-q"], ["add", "."], ["commit", "-q", "-m", subject]): + subprocess.run(["git", *args], cwd=root, env=env, check=True, capture_output=True) + return root + + +@pytest.fixture(name="captured") +def _captured() -> Tuple[str, List[bytes]]: + bodies: List[bytes] = [] + lock = threading.Lock() + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + raw = self.rfile.read(int(self.headers.get("content-length") or 0)) + if "count_tokens" not in self.path: + with lock: + bodies.append(raw) + payload = json.dumps({"input_tokens": 10} if "count_tokens" in self.path else _STUB_REPLY).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *_args: object) -> None: + return + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}", bodies + finally: + server.shutdown() + + +def test_two_builds_send_the_same_request_bytes(captured: Tuple[str, List[bytes]], tmp_path: Path) -> None: + base_url, bodies = captured + limiter = RateLimiter(state_dir=tmp_path / "limiter") + checkouts = (_make_repo(tmp_path / "build-1", "first"), _make_repo(tmp_path / "build-2", "second")) + origin = Path.cwd() + + sent = [] + for checkout in checkouts: + shutil.rmtree(Path(_stable_cli_state()[0]).parent, ignore_errors=True) + os.chdir(checkout) + try: + before = len(bodies) + run_claude( + prompt="say ok", + model="claude-haiku-4-5", + base_url=base_url, + api_key="stub", + extra_env={"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"}, + rate_limiter=limiter, + ) + sent.append(bodies[before:]) + finally: + os.chdir(origin) + + assert sent[0], "the CLI sent no request to the stub, so there is nothing to compare" + assert sent[0] == sent[1] + + +def test_concurrent_cells_do_not_collide_on_the_pinned_session( + captured: Tuple[str, List[bytes]], tmp_path: Path +) -> None: + base_url, bodies = captured + limiter = RateLimiter(state_dir=tmp_path / "limiter") + + def one(_index: int) -> int: + return run_claude( + prompt="say ok", + model="claude-haiku-4-5", + base_url=base_url, + api_key="stub", + extra_env={"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1"}, + rate_limiter=limiter, + ).exit_code + + with ThreadPoolExecutor(max_workers=4) as pool: + codes = list(pool.map(one, range(4))) + + assert codes == [0, 0, 0, 0] + assert bodies, "the CLI sent no request to the stub, so there is nothing to compare" + assert set(Counter(bodies).values()) == {4} diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index 447e8cc0bbb..3fad87c7479 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -132,6 +132,57 @@ def _make_isolated_home() -> str: return tempfile.mkdtemp(prefix="claude-cli-home-") +_FIXED_CLI_USER_ID = "0" * 64 +_FIXED_CLI_SESSION_ID = "00000000-0000-4000-8000-000000000000" + + +def _seed_cli_identity(config_dir: str) -> None: + """Pin the device id the CLI would otherwise mint per config directory. + + It mints 32 random bytes on first run, writes them to `.claude.json` as + `userID`, and sends them in `metadata.user_id` forever after, so the value + is stable for exactly as long as that file lives. Pinning it, and the + session id passed beside it, costs nothing: both feed abuse detection + rather than quota, caching or continuity.""" + path = os.path.join(config_dir, ".claude.json") + try: + with open(path, encoding="utf-8") as handle: + if json.load(handle).get("userID") == _FIXED_CLI_USER_ID: + return + except (OSError, ValueError): + pass + staged = f"{path}.{os.getpid()}" + with open(staged, "w", encoding="utf-8") as handle: + json.dump({"userID": _FIXED_CLI_USER_ID}, handle) + os.replace(staged, path) + + +def _stable_cli_state() -> Tuple[str, str]: + """Config directory and working directory for the CLI, at fixed paths. + + Both reach the request body. The memory directory the system prompt + names is `$CLAUDE_CONFIG_DIR/projects//memory`, and a working + directory inside a git repository also contributes its branch and recent + commits. So a per-invocation config directory rewrites every body, and + inheriting the checkout rewrites every body once per candidate, which is + why the shared provider cache could never serve a Claude Code cell. + Pinning both makes the bodies repeatable across builds. + + This narrows what survives rather than widening it: HOME stays fresh and + empty per invocation, so the isolation `_make_isolated_home` describes is + unchanged, and the CLI's own state no longer outlives the pod either. The + working directory is deliberately not the checkout, so a model-directed + `Read` sees an empty directory instead of the repository. + """ + root = os.path.join(tempfile.gettempdir(), f"litellm-e2e-claude-{os.getuid()}") + config_dir = os.path.join(root, "config") + workspace = os.path.join(root, "workspace") + for path in (root, config_dir, workspace): + os.makedirs(path, mode=0o700, exist_ok=True) + _seed_cli_identity(config_dir) + return config_dir, workspace + + class ClaudeCLIError(RuntimeError): """Raised when the `claude` CLI cannot be invoked or returns a fatal error.""" @@ -222,6 +273,9 @@ def run_claude( "--verbose", "--model", model, + "--session-id", + _FIXED_CLI_SESSION_ID, + "--no-session-persistence", ] if extra_args: cmd.extend(extra_args) @@ -244,6 +298,8 @@ def run_claude( # regardless of how the subprocess exits. isolated_home = _make_isolated_home() env["HOME"] = isolated_home + config_dir, workspace = _stable_cli_state() + env["CLAUDE_CONFIG_DIR"] = config_dir if extra_env: env.update(extra_env) @@ -262,6 +318,7 @@ def run_claude( completed = run_fn( cmd, env=env, + cwd=workspace, input=stdin_input, capture_output=True, text=True, From 2481146727fe3b2613df96ccf8f568f5b0a3cc72 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 12:44:53 -0700 Subject: [PATCH 113/187] docs(e2e): say why the CLI-driving cells needed a driver fix, not a rule --- tests/e2e/PROVIDER_CACHE.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/PROVIDER_CACHE.md b/tests/e2e/PROVIDER_CACHE.md index e06b3c01653..d37b37eeba7 100644 --- a/tests/e2e/PROVIDER_CACHE.md +++ b/tests/e2e/PROVIDER_CACHE.md @@ -16,6 +16,8 @@ Requests that differ only by their markers therefore share a canonical identity, Two different tests never share a recording, and a provider call made outside any test (fixtures, session setup) is never cached, because the identity has no test node id to bind to +A client that varies its own request between runs defeats that identity without breaking any rule, and the Claude Code compat cells did. The CLI sends a device id and a session id in `metadata.user_id`, and its system prompt names both its memory directory and its working directory, adding the branch and recent commits when that directory is a git repository. Driven with a fresh HOME and the checkout as its working directory, every cell sent different bytes every build. The fix belongs in the driver rather than here: `claude_code/cli_driver.py` pins the config directory, the working directory and both identifiers, which is why the cache needs no rule for any of it. Normalizing them instead would have hidden a real defect class, since a rule cannot tell a client's own churn from a value a test means to assert on + Provider `Set-Cookie` headers are dropped before validation and never recorded: the edge already withholds them from the proxy, and OpenAI responses always carry Cloudflare bot-management cookies An eligible miss calls the provider. A complete successful response is stored immediately even if a later test assertion fails. Provider errors, malformed responses, truncated streams and cancelled captures are not stored. Cache reads, writes and lease failures fall through to normal provider behavior; they introduce no provider retry. An already-started response cannot be restarted after a delivery failure From 7ff8945da411ed30891b17dfbe707c49b4a76803 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 19:45:48 +0000 Subject: [PATCH 114/187] fix(router): sync both usage keys from Redis even when one increment is zero A stream counted before its usage is known increments TPM by zero, so the worker that served it never refreshed its local TPM value from Redis and the first byte headers reported the token count another worker had already consumed. Both pipeline operations now always run, matching the pre-change callback, so the returned values refresh both worker local keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 1 - .../router_unit_tests/test_router_helper_utils.py | 15 +++++++++++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index cc11a21d966..6e285d54550 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8281,7 +8281,6 @@ class Router: pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [ RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value) for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment)) - if increment_value > 0 ] post_increment_values: Final = await self.cache.async_increment_cache_pipeline( increment_list=pipeline_operations, diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 664f15889d3..4a02ec62064 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1041,7 +1041,7 @@ async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_ headers = _ratelimit_headers(stream) assert headers["x-ratelimit-remaining-tokens"] == 1000 assert headers["x-ratelimit-remaining-requests"] == 99 - assert await router.get_model_group_usage("gpt-5-mini") == (None, 1) + assert await router.get_model_group_usage("gpt-5-mini") == (0, 1) chunks = [chunk async for chunk in stream] total_tokens = chunks[-1].usage.total_tokens @@ -1078,7 +1078,7 @@ async def test_deployment_callback_on_success_adds_only_uncounted_tokens(): ) assert tpm_key is not None - assert await router.get_model_group_usage("gpt-5-mini") == (40, None) + assert await router.get_model_group_usage("gpt-5-mini") == (40, 0) class _GatedIncrementCache(DualCache): @@ -1228,6 +1228,17 @@ async def test_headers_on_fresh_worker_reflect_shared_redis_usage(): assert headers["x-ratelimit-remaining-requests"] == 96 assert headers["x-ratelimit-remaining-tokens"] == 1000 - tokens_on_a - response.usage.total_tokens + counted_tokens = tokens_on_a + response.usage.total_tokens + for _ in range(2): + response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") + counted_tokens += response.usage.total_tokens + + stream = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong", stream=True) + stream_headers = _ratelimit_headers(stream) + assert stream_headers["x-ratelimit-remaining-requests"] == 93 + assert stream_headers["x-ratelimit-remaining-tokens"] == 1000 - counted_tokens + assert [chunk async for chunk in stream] + @pytest.mark.asyncio async def test_get_model_group_io_token_usage_sums_across_deployments(): From 7679e42736671a9e65b46d9fcb32d77e3f59170e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 12:48:30 -0700 Subject: [PATCH 115/187] test: cover provider wire contracts, streaming and recovery Adds owned local TCP peers, a Redis process helper and SigV4 helpers to tests/integration, and integration contracts for Anthropic and Bedrock wire shapes, Bedrock role configuration, S3 wire, observed routing, Redis recovery and stream contracts. Consolidates the four commits previously stacked on litellm_integration_accounting onto its main-merged tip --- tests/integration/README.md | 4 + tests/integration/_support/process.py | 6 +- tests/integration/_support/redis_process.py | 116 ++++++++++++++ tests/integration/_support/sigv4.py | 25 +++ tests/integration/_support/wire.py | 113 +++++++++++++ tests/integration/contracts.json | 47 ++++++ .../providers/test_anthropic_wire.py | 61 +++++++ .../providers/test_bedrock_auth_wire.py | 99 ++++++++++++ .../test_bedrock_role_configuration.py | 75 +++++++++ tests/integration/providers/test_s3_wire.py | 111 +++++++++++++ .../routing/test_observed_routing.py | 98 ++++++++++++ .../routing/test_redis_recovery.py | 59 +++++++ .../streaming/test_stream_contracts.py | 149 ++++++++++++++++++ 13 files changed, 960 insertions(+), 3 deletions(-) create mode 100644 tests/integration/_support/redis_process.py create mode 100644 tests/integration/_support/sigv4.py create mode 100644 tests/integration/_support/wire.py create mode 100644 tests/integration/providers/test_anthropic_wire.py create mode 100644 tests/integration/providers/test_bedrock_auth_wire.py create mode 100644 tests/integration/providers/test_bedrock_role_configuration.py create mode 100644 tests/integration/providers/test_s3_wire.py create mode 100644 tests/integration/routing/test_observed_routing.py create mode 100644 tests/integration/routing/test_redis_recovery.py create mode 100644 tests/integration/streaming/test_stream_contracts.py diff --git a/tests/integration/README.md b/tests/integration/README.md index 79f25f760f9..64c0d7c9412 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -21,3 +21,7 @@ Fixtures must contain synthetic data only. Keep private incident records and sou Database cases own their temporary schemas, roles, constraints and proxy processes. They prove reader-versus-writer execution with PostgreSQL lock observations, exercise real transaction wait limits and verify rollback after a reached database failure Accounting cases compare persisted input and output cost components against literal rates, including zero and default prices. Cache state models assert actual upstream calls, response identity and every persisted charge. Generated accounting tests have a 180-second test limit to accommodate the asynchronous spend writer; CircleCI keeps the whole shard capped at 11 minutes + +Provider contracts exercise actual TCP requests with synthetic credentials and local protocol peers. The S3 verifier uses independently implemented equations, a published known-answer vector, a fixed signing clock and deliberately invalid signed requests. Bedrock cases clear ambient AWS credential sources and check the literal model path, loaded role references, STS requests and bearer-only behavior + +Streaming checks send real HTTP transfer chunks, including one-byte partitions, fragmented tools, incomplete transfers and a cancellation barrier. They assert meaningful text, tool arguments, final usage and persisted cost. The Redis recovery case owns a separate database and Redis process, uses the supported one-second circuit-breaker recovery setting, waits for the real subscriber and verifies response data in Redis after restart. CircleCI reuses its existing Redis image for that extra process; it never pulls an image during tests diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 0d66ecc9d90..84ee2ad1b79 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -46,13 +46,13 @@ def stop_root_process(process: subprocess.Popen[bytes]) -> bool: @contextmanager -def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str]) -> Iterator[Gateway]: +def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str], *, config: Path | None = None, remove_environment: tuple[str, ...] = ()) -> Iterator[Gateway]: with socket.socket() as reserve: reserve.bind(("127.0.0.1", 0)) port: Final = reserve.getsockname()[1] root: Final = Path(__file__).resolve().parents[3] environment: Final = { - **os.environ, + **{name: value for name, value in os.environ.items() if name not in remove_environment}, "LITELLM_MASTER_KEY": gateway.key, "LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"), "STORE_MODEL_IN_DB": "True", @@ -67,7 +67,7 @@ def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str]) "-m", "integration._support.proxy", "--config", - "tests/integration/proxy_config.yaml", + str(config or "tests/integration/proxy_config.yaml"), "--host", "127.0.0.1", "--port", diff --git a/tests/integration/_support/redis_process.py b/tests/integration/_support/redis_process.py new file mode 100644 index 00000000000..86abcbe024e --- /dev/null +++ b/tests/integration/_support/redis_process.py @@ -0,0 +1,116 @@ +import os +import shutil +import signal +import socket +import subprocess +import time +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final, TextIO + +from redis import Redis +from redis.exceptions import ConnectionError as RedisConnectionError + + +@dataclass +class OwnedRedis: + host: str + port: int + command: tuple[str, ...] + log: TextIO + pid_file: str + process: subprocess.Popen | None = None + server_pid: int | None = None + + def start(self) -> None: + assert self.process is None + self.process = subprocess.Popen(self.command, stdout=self.log, stderr=subprocess.STDOUT, start_new_session=True) + deadline: Final = time.monotonic() + 8 + with Redis(host=self.host, port=self.port, socket_connect_timeout=0.2, socket_timeout=0.2) as client: + while True: + assert self.process.poll() is None, "Owned Redis exited before readiness" + try: + if client.ping(): + actual: Final = int(client.info("server")["process_id"]) + expected: Final = self.process.pid if self.command[0] != "docker" else int(subprocess.check_output(["docker", "exec", "redis-cache", "cat", self.pid_file], timeout=2)) + assert actual == expected, "Redis readiness reached a different process" + self.server_pid = actual + return + except RedisConnectionError: + pass + assert time.monotonic() < deadline, "Owned Redis readiness deadline exceeded" + time.sleep(0.05) + + def stop(self) -> None: + assert self.process is not None + failure = None + forced = False + try: + if self.process.poll() is None: + with Redis(host=self.host, port=self.port, socket_connect_timeout=1, socket_timeout=1) as client: + assert int(client.info("server")["process_id"]) == self.server_pid, "Redis ownership changed before shutdown" + client.shutdown(nosave=True) + except Exception as error: + failure = error + finally: + try: + self.process.wait(timeout=5) + except subprocess.TimeoutExpired: + forced = True + self.signal(signal.SIGTERM) + try: + self.process.wait(timeout=5) + except subprocess.TimeoutExpired: + self.signal(signal.SIGKILL) + self.process.wait(timeout=3) + self.process = None + self.server_pid = None + with Redis(host=self.host, port=self.port, socket_connect_timeout=0.2, socket_timeout=0.2) as client: + try: + client.ping() + except RedisConnectionError: + stopped = True + else: + stopped = False + assert stopped, "Owned Redis still serves after shutdown" + assert failure is None and not forced, f"Owned Redis required shutdown recovery: {failure!r}" + + def signal(self, action: signal.Signals) -> None: + assert self.process is not None + if self.command[0] != "docker": + self.process.send_signal(action) + return + pid: Final = int(subprocess.check_output(["docker", "exec", "redis-cache", "cat", self.pid_file], timeout=2)) + command: Final = subprocess.check_output(["docker", "exec", "redis-cache", "cat", f"/proc/{pid}/cmdline"], timeout=2) + assert self.pid_file.encode() in command, "Redis process ownership changed" + subprocess.run(["docker", "exec", "redis-cache", "kill", f"-{int(action)}", str(pid)], check=True, timeout=2) + + +@contextmanager +def owned_redis(directory: Path) -> Iterator[OwnedRedis]: + binary: Final = shutil.which("redis-server") + if binary: + with socket.socket() as reservation: + reservation.bind(("127.0.0.1", 0)) + port = reservation.getsockname()[1] + host = "127.0.0.1" + prefix = (binary,) + else: + host = subprocess.check_output(["docker", "inspect", "--format", "{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}", "redis-cache"], text=True).strip() + assert host, "CircleCI owned Redis container has no address" + port = 16379 + prefix = ("docker", "exec", "redis-cache", "redis-server") + output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory))) + output.mkdir(parents=True, exist_ok=True) + with (output / "owned-redis-recovery.log").open("w") as log: + pid_file: Final = str(directory / "owned-redis.pid") if binary else f"/tmp/integration-redis-{uuid.uuid4().hex}.pid" + server: Final = OwnedRedis(host, port, (*prefix, "--port", str(port), "--set-proc-title", "no", "--pidfile", pid_file, "--bind", "0.0.0.0" if not binary else "127.0.0.1", "--protected-mode", "no", "--save", "", "--appendonly", "no"), log, pid_file) + try: + server.start() + yield server + finally: + if server.process is not None: + server.stop() diff --git a/tests/integration/_support/sigv4.py b/tests/integration/_support/sigv4.py new file mode 100644 index 00000000000..e02283a719d --- /dev/null +++ b/tests/integration/_support/sigv4.py @@ -0,0 +1,25 @@ +import hashlib +import hmac +from collections.abc import Mapping +from typing import Final + + +def encoded_path(value: str) -> str: + safe: Final = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_.~/" + return "".join(chr(byte) if byte in safe else f"%{byte:02X}" for byte in value.encode("utf-8")) + + +def signature( + method: str, path: str, headers: Mapping[str, str], signed: str, body: bytes, secret: str, scope: str, +) -> tuple[str, str]: + """AWS SigV4 equations, independent of botocore and LiteLLM's signer.""" + canonical_headers: Final = "".join(name + ":" + " ".join(headers[name].split()) + "\n" for name in signed.split(";")) + canonical: Final = "\n".join((method, path, "", canonical_headers, signed, hashlib.sha256(body).hexdigest())) + canonical_hash: Final = hashlib.sha256(canonical.encode()).hexdigest() + date, region, service, terminator = scope.split("/") + assert terminator == "aws4_request" + key = ("AWS4" + secret).encode() + for part in (date, region, service, terminator): + key = hmac.new(key, part.encode(), hashlib.sha256).digest() + to_sign: Final = "\n".join(("AWS4-HMAC-SHA256", headers["x-amz-date"], scope, canonical_hash)) + return canonical_hash, hmac.new(key, to_sign.encode(), hashlib.sha256).hexdigest() diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py new file mode 100644 index 00000000000..acc51dd4497 --- /dev/null +++ b/tests/integration/_support/wire.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import threading +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from queue import SimpleQueue +from typing import Final + + +@dataclass(frozen=True, slots=True) +class Request: + method: str + target: str + headers: Mapping[str, str] + body: bytes + + +@dataclass(frozen=True, slots=True) +class Reply: + status: int = 200 + body: bytes = b"{}" + content_type: str = "application/json" + chunks: tuple[bytes, ...] | None = None + abort_after: int | None = None + gate_after_first: threading.Event | None = None + + +@dataclass(frozen=True, slots=True) +class Wire: + url: str + received: SimpleQueue[Request] + disconnected: SimpleQueue[str] + + def drain(self) -> tuple[Request, ...]: + return tuple(self.received.get_nowait() for _ in range(self.received.qsize())) + + +@contextmanager +def wire_server(respond: Callable[[Request], Reply]) -> Iterator[Wire]: + """Owned TCP peer; requests traverse the real HTTP client and serialization.""" + received: Final[SimpleQueue[Request]] = SimpleQueue() + errors: Final[SimpleQueue[Exception]] = SimpleQueue() + disconnected: Final[SimpleQueue[str]] = SimpleQueue() + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + timeout = 5 + + def respond(self) -> None: + request: Final = Request( + self.command, self.path, + {name.lower(): value for name, value in self.headers.items()}, + self.rfile.read(int(self.headers.get("content-length", "0"))), + ) + received.put(request) + try: + reply = respond(request) + except Exception as error: + errors.put(error) + reply = Reply(status=500) + self.send_response(reply.status) + self.send_header("content-type", reply.content_type) + if reply.chunks is None: + self.send_header("content-length", str(len(reply.body))) + else: + self.send_header("transfer-encoding", "chunked") + self.send_header("connection", "close") + self.end_headers() + try: + if reply.chunks is None: + self.wfile.write(reply.body) + else: + for index, chunk in enumerate(reply.chunks): + if reply.abort_after == index: + break + self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk)) + self.wfile.flush() + if index == 0 and reply.gate_after_first is not None: + assert reply.gate_after_first.wait(timeout=5), "Stream barrier was never released" + else: + self.wfile.write(b"0\r\n\r\n") + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + disconnected.put(request.target) + except Exception as error: + errors.put(error) + self.close_connection = True + + do_POST = respond + do_PUT = respond + do_GET = respond + do_DELETE = respond + + def log_message(self, format: str, *args: object) -> None: + pass + + class OwnedHTTPServer(ThreadingHTTPServer): + daemon_threads = False + + with OwnedHTTPServer(("127.0.0.1", 0), Handler) as server: + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread.start() + try: + yield Wire(f"http://127.0.0.1:{server.server_port}", received, disconnected) + finally: + server.shutdown() + thread.join(timeout=6) + assert not thread.is_alive(), "Owned HTTP server survived cleanup" + server.server_close() + failure: Final = None if errors.empty() else errors.get_nowait() + assert failure is None, f"Owned HTTP peer failed: {failure!r}" diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 8c09048c243..67b45cbbc54 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -103,6 +103,53 @@ ], "tests/integration/spend/test_cache_and_quota.py::test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows": [ "quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge" + ], + "tests/integration/providers/test_s3_wire.py::test_sigv4_verifier_matches_published_put_and_rejects_corruption": [ + "other.provider_wire.s3.verifier_known_answer_and_negative_controls" + ], + "tests/integration/providers/test_s3_wire.py::test_s3_sync_and_async_uploads_pass_independent_wire_verification": [ + "other.provider_wire.s3.sync_async_reserved_keys_are_signed_and_accepted" + ], + "tests/integration/providers/test_bedrock_auth_wire.py::test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials": [ + "other.provider_wire.bedrock.bearer_sdk_skips_credential_chain" + ], + "tests/integration/providers/test_bedrock_auth_wire.py::test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload": [ + "other.provider_wire.bedrock.bearer_db_yaml_survives_reload" + ], + "tests/integration/streaming/test_stream_contracts.py::test_generated_tcp_partitions_preserve_unicode_text_identity_and_final_usage": [ + "other.streaming.byte_partitions.preserve_text_identity_and_usage" + ], + "tests/integration/streaming/test_stream_contracts.py::test_fragmented_tool_names_and_arguments_keep_each_call_identity": [ + "other.streaming.tools.fragmented_calls_keep_independent_arguments" + ], + "tests/integration/streaming/test_stream_contracts.py::test_proxy_stream_usage_visibility_keeps_exact_persisted_charge": [ + "other.streaming.usage.client_visibility_preserves_persisted_accounting" + ], + "tests/integration/streaming/test_stream_contracts.py::test_truncated_http_stream_is_an_error_and_next_stream_succeeds": [ + "other.streaming.failure.truncated_transport_raises_and_control_recovers" + ], + "tests/integration/streaming/test_stream_contracts.py::test_client_cancellation_releases_the_actual_provider_connection": [ + "other.streaming.cancellation.closes_actual_provider_connection" + ], + "tests/integration/routing/test_observed_routing.py::test_retry_counts_and_public_errors_match_actual_provider_attempts": [ + "other.routing.retries.several_attempts_reach_success_without_hidden_retries", + "other.routing.errors.nonretryable_and_exhausted_failures_remain_errors" + ], + "tests/integration/routing/test_observed_routing.py::test_loaded_fallback_selects_expected_deployment_and_keeps_response_identity": [ + "other.routing.fallback.loaded_configuration_selects_only_permitted_target" + ], + "tests/integration/routing/test_observed_routing.py::test_saved_deployment_target_update_changes_wire_and_preserves_control": [ + "other.routing.alias_update.persisted_target_changes_only_selected_route" + ], + "tests/integration/providers/test_bedrock_role_configuration.py::test_role_reference_from_db_and_yaml_reaches_real_sts_http_and_bedrock": [ + "other.provider_wire.bedrock.db_yaml_role_reference_reaches_sts_and_signed_request" + ], + "tests/integration/routing/test_redis_recovery.py::test_owned_redis_outage_recovers_requests_and_real_response_cache": [ + "other.routing.redis.owned_outage_recovers_serving_and_response_cache" + ], + "tests/integration/providers/test_anthropic_wire.py::test_anthropic_tool_history_and_cache_tokens_keep_wire_and_accounting_contracts": [ + "other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields", + "quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates" ] } } diff --git a/tests/integration/providers/test_anthropic_wire.py b/tests/integration/providers/test_anthropic_wire.py new file mode 100644 index 00000000000..64160fa85aa --- /dev/null +++ b/tests/integration/providers/test_anthropic_wire.py @@ -0,0 +1,61 @@ +import json +import uuid +from typing import Final + +import pytest + +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server + + +@pytest.mark.covers("other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields", "quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates") +def test_anthropic_tool_history_and_cache_tokens_keep_wire_and_accounting_contracts(gateway: Gateway) -> None: + identity: Final = "anthropic-wire-" + uuid.uuid4().hex + tool_schema: Final = {"type": "object", "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}}, "required": ["x", "y"]} + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages" + assert request.headers["x-api-key"] == "synthetic-anthropic-key" + body: Final = json.loads(request.body) + assert body["model"] == "claude-sonnet-4-5-20250929" + assert body["system"] == [{"type": "text", "text": "synthetic policy", "cache_control": {"type": "ephemeral"}}] + assert body["tools"][0]["name"] == "add" and body["tools"][0]["input_schema"] == tool_schema + assert body["max_tokens"] == 16 + assert not {"timeout", "stream_chunk_size", "litellm_params", "litellm_metadata", "rpm", "tpm"}.intersection(body) + messages: Final = body["messages"] + assert [message["role"] for message in messages] == ["user", "assistant", "user"] + assert messages[0]["content"] == [{"type": "text", "text": "first"}] + assert messages[1]["content"] == [{"type": "tool_use", "id": "history-call", "name": "add", "input": {"x": 1, "y": 2}}] + assert messages[2]["content"] == [{"type": "tool_result", "tool_use_id": "history-call", "content": "3"}, {"type": "text", "text": "next"}] + return Reply(body=json.dumps({"id": identity, "type": "message", "role": "assistant", "model": "claude-sonnet-4-5-20250929", "content": [{"type": "tool_use", "id": "next-call", "name": "add", "input": {"x": 3, "y": 4}}], "stop_reason": "tool_use", "stop_sequence": None, "usage": {"input_tokens": 10, "output_tokens": 4, "cache_read_input_tokens": 5, "cache_creation_input_tokens": 7}}).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="anthropic/claude-sonnet-4-5-20250929", api_base=wire.url, api_key="synthetic-anthropic-key", input_cost_per_token=0.001, output_cost_per_token=0.002, cache_read_input_token_cost=0.0001, cache_creation_input_token_cost=0.002) + response: Final = gateway.request("POST", "/v1/chat/completions", { + "model": model, "max_tokens": 16, "timeout": 5, + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "synthetic policy", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": "first"}, + {"role": "assistant", "tool_calls": [{"id": "history-call", "type": "function", "function": {"name": "add", "arguments": '{"x":1,"y":2}'}}]}, + {"role": "tool", "tool_call_id": "history-call", "content": "3"}, + {"role": "user", "content": "next"}, + ], + "tools": [{"type": "function", "function": {"name": "add", "parameters": tool_schema}}], + }) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["id"].startswith("chatcmpl-") + assert body["choices"][0]["finish_reason"] == "tool_calls" + tool: Final = body["choices"][0]["message"]["tool_calls"][0] + assert tool["id"] == "next-call" and tool["function"]["name"] == "add" + assert json.loads(tool["function"]["arguments"]) == {"x": 3, "y": 4} + assert body["usage"]["prompt_tokens"] == 22 and body["usage"]["completion_tokens"] == 4 + assert len(wire.drain()) == 1 + rows: Final = eventually(lambda: read_rows('SELECT spend, prompt_tokens, completion_tokens, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (body["id"],)), lambda values: len(values) == 1, seconds=70) + assert float(rows[0]["spend"]) == pytest.approx(10 * 0.001 + 5 * 0.0001 + 7 * 0.002 + 4 * 0.002) + assert rows[0]["prompt_tokens"] == 22 and rows[0]["completion_tokens"] == 4 + metadata: Final = rows[0]["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + assert parsed["cost_breakdown"]["input_cost"] == pytest.approx(0.0245) + assert parsed["cost_breakdown"]["output_cost"] == pytest.approx(0.008) diff --git a/tests/integration/providers/test_bedrock_auth_wire.py b/tests/integration/providers/test_bedrock_auth_wire.py new file mode 100644 index 00000000000..bd24dc171ba --- /dev/null +++ b/tests/integration/providers/test_bedrock_auth_wire.py @@ -0,0 +1,99 @@ +import asyncio +import json +import os +import uuid +from pathlib import Path +from typing import Final + +import pytest +import yaml + +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +MODEL: Final = "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0" +TOKEN: Final = "synthetic-bedrock-bearer" +RESPONSE: Final = json.dumps({ + "output": {"message": {"role": "assistant", "content": [{"text": "bedrock wire control"}]}}, + "stopReason": "end_turn", "usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}, + "metrics": {"latencyMs": 1}, +}).encode() + + +def bearer_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse" + assert request.headers["authorization"] == f"Bearer {TOKEN}" + assert "x-amz-security-token" not in request.headers + body: Final = json.loads(request.body) + assert body["messages"] == [{"role": "user", "content": [{"text": "synthetic bearer request"}]}] + assert body["system"] == [{"text": "synthetic system"}] + assert body["inferenceConfig"]["maxTokens"] == 16 + assert not {"timeout", "stream_chunk_size", "litellm_params", "litellm_metadata", "api_key"}.intersection(body) + return Reply(body=RESPONSE) + + +@pytest.mark.covers("other.provider_wire.bedrock.bearer_sdk_skips_credential_chain") +async def test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + import litellm + + empty: Final = tmp_path / "empty-aws-config" + empty.write_text("") + for name in tuple(name for name in os.environ if name.startswith("AWS_")): + monkeypatch.delenv(name, raising=False) + for name, value in {"AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}.items(): + monkeypatch.setenv(name, value) + with wire_server(bearer_peer) as wire: + with pytest.raises(litellm.APIConnectionError, match=r"config profile .* could not be found"): + await asyncio.to_thread(litellm.completion, model=MODEL, aws_profile_name="integration-profile-must-not-be-read", aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url, messages=[{"role": "user", "content": "synthetic credential control"}], timeout=5, num_retries=0) + assert wire.drain() == () + for source in ("argument", "environment"): + if source == "environment": + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", TOKEN) + parameters: Final = { + "model": MODEL, "api_key": TOKEN if source == "argument" else None, + "aws_region_name": "us-east-1", "aws_profile_name": "integration-profile-must-not-be-read", + "aws_bedrock_runtime_endpoint": wire.url, "timeout": 5, "num_retries": 0, + "messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}], + "max_tokens": 16, + } + for asynchronous in (False, True): + result: Final = await litellm.acompletion(**parameters) if asynchronous else await asyncio.to_thread(litellm.completion, **parameters) + assert result.choices[0].message.content == "bedrock wire control" + assert result.choices[0].finish_reason == "stop" + assert result.usage.prompt_tokens == 11 and result.usage.completion_tokens == 4 + assert len(wire.drain()) == 1 + + +@pytest.mark.covers("other.provider_wire.bedrock.bearer_db_yaml_survives_reload") +def test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload(gateway: Gateway, tmp_path: Path) -> None: + empty: Final = tmp_path / "empty-aws-config" + empty.write_text("") + with wire_server(bearer_peer) as wire: + parameters: Final = { + "model": MODEL, "api_key": "os.environ/INTEGRATION_BEARER_TOKEN", "aws_region_name": "us-east-1", + "aws_profile_name": "integration-profile-must-not-be-read", "aws_bedrock_runtime_endpoint": wire.url, + } + alias: Final = f"integration-yaml-{uuid.uuid4().hex}" + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["model_list"] = [{"model_name": alias, "litellm_params": parameters, "model_info": {"id": alias}}] + path: Final = tmp_path / "bedrock.yaml" + path.write_text(yaml.safe_dump(configuration)) + overrides: Final = {"INTEGRATION_BEARER_TOKEN": TOKEN, "AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"} + with owned_proxy(gateway, tmp_path, overrides, config=path, remove_environment=tuple(name for name in os.environ if name.startswith("AWS_"))) as candidate, candidate.scenario() as scenario: + database_model: Final = scenario.model(**parameters) + for generation in range(2): + for model in (alias, database_model): + response: Final = candidate.request("POST", "/v1/chat/completions", { + "model": model, "messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}], + "max_tokens": 16, "cache": {"no-cache": True}, + }) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control" + assert response.json()["usage"]["total_tokens"] == 15 + assert len(wire.drain()) == 1, f"Expected actual provider call after reload {generation}" + if generation == 0: + entries: Final = candidate.get("/model/info")["data"] + target: Final = next(entry for entry in entries if entry["model_name"] == database_model) + response: Final = candidate.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "bearer reload"}}) + assert response.status_code == 200, response.text diff --git a/tests/integration/providers/test_bedrock_role_configuration.py b/tests/integration/providers/test_bedrock_role_configuration.py new file mode 100644 index 00000000000..ac8edbdfde0 --- /dev/null +++ b/tests/integration/providers/test_bedrock_role_configuration.py @@ -0,0 +1,75 @@ +import json +import os +import uuid +from pathlib import Path +from typing import Final +from urllib.parse import parse_qs + +import pytest +import yaml + +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from integration.providers.test_bedrock_auth_wire import MODEL, RESPONSE + + +@pytest.mark.covers("other.provider_wire.bedrock.db_yaml_role_reference_reaches_sts_and_signed_request") +def test_role_reference_from_db_and_yaml_reaches_real_sts_http_and_bedrock(gateway: Gateway, tmp_path: Path) -> None: + role: Final = "arn:aws:iam::123456789012:role/integration-" + uuid.uuid4().hex + assumed_key: Final = "ASIAINTEGRATION000001" + assumed_token: Final = "synthetic-assumed-session-token" + + def sts(request: Request) -> Reply: + parameters: Final = parse_qs(request.body.decode()) + action: Final = parameters["Action"][0] + assert request.method == "POST" and action in {"GetCallerIdentity", "AssumeRole"} + if action == "GetCallerIdentity": + result = "arn:aws:iam::123456789012:user/integration-sourceintegration-source123456789012" + else: + assert parameters["RoleArn"] == [role] + assert parameters["RoleSessionName"][0] in {"integration-yaml-session", "integration-db-session"} + result = f"{assumed_key}synthetic-assumed-secret-key-for-testing{assumed_token}2035-01-01T00:00:00Zarn:aws:sts::123456789012:assumed-role/integration/sessionintegration:session0" + return Reply(content_type="text/xml", body=f'<{action}Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">{result}synthetic-sts-request'.encode()) + + def bedrock(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse" + assert f"Credential={assumed_key}/" in request.headers["authorization"] + assert request.headers["x-amz-security-token"] == assumed_token + assert json.loads(request.body)["messages"][0]["content"][0]["text"] == "synthetic role request" + return Reply(body=RESPONSE) + + with wire_server(sts) as authority, wire_server(bedrock) as provider: + parameters: Final = { + "model": MODEL, "aws_region_name": "us-east-1", "aws_role_name": "os.environ/INTEGRATION_ROLE_ARN", + "aws_session_name": "integration-yaml-session", "aws_bedrock_runtime_endpoint": provider.url, + "aws_sts_endpoint": authority.url, + } + alias: Final = "integration-role-yaml-" + uuid.uuid4().hex + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["model_list"] = [{"model_name": alias, "litellm_params": parameters, "model_info": {"id": alias}}] + path: Final = tmp_path / "roles.yaml" + path.write_text(yaml.safe_dump(configuration)) + empty: Final = tmp_path / "empty-aws-config" + empty.write_text("") + overrides: Final = { + "INTEGRATION_ROLE_ARN": role, "AWS_ACCESS_KEY_ID": "AKIAINTEGRATION000001", "AWS_SECRET_ACCESS_KEY": "synthetic-source-secret-key-for-testing", + "AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", + "AWS_ENDPOINT_URL_STS": authority.url, "AWS_DEFAULT_REGION": "us-east-1", "LITELLM_RUST": "false", + } + with owned_proxy(gateway, tmp_path, overrides, config=path, remove_environment=tuple(name for name in os.environ if name.startswith("AWS_"))) as candidate, candidate.scenario() as scenario: + database_model: Final = scenario.model(**{**parameters, "api_key": None, "aws_session_name": "integration-db-session"}) + for generation in range(2): + for model in (alias, database_model): + response: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "synthetic role request"}], "cache": {"no-cache": True}}) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control" + assert response.json()["usage"]["total_tokens"] == 15 + assert len(provider.drain()) == 1 + if generation == 0: + target: Final = next(entry for entry in candidate.get("/model/info")["data"] if entry["model_name"] == database_model) + response: Final = candidate.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "role reload"}}) + assert response.status_code == 200, response.text + assumed: Final = tuple(parse_qs(request.body.decode()) for request in authority.drain() if parse_qs(request.body.decode())["Action"] == ["AssumeRole"]) + assert {entry["RoleSessionName"][0] for entry in assumed} == {"integration-yaml-session", "integration-db-session"} + assert all(entry["RoleArn"] == [role] for entry in assumed) diff --git a/tests/integration/providers/test_s3_wire.py b/tests/integration/providers/test_s3_wire.py new file mode 100644 index 00000000000..e6c5ac18a49 --- /dev/null +++ b/tests/integration/providers/test_s3_wire.py @@ -0,0 +1,111 @@ +import asyncio +import base64 +import hashlib +import hmac +import json +from datetime import datetime +from typing import Final + +import httpx +import pytest + +from integration._support.sigv4 import encoded_path, signature +from integration._support.wire import Reply, Request, wire_server + +ACCESS: Final = "AKIAIOSFODNN7EXAMPLE" +SECRET: Final = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + + +@pytest.mark.covers("other.provider_wire.s3.verifier_known_answer_and_negative_controls") +def test_sigv4_verifier_matches_published_put_and_rejects_corruption() -> None: + # Public AWS example credentials and PUT vector, not an active account: + # https://docs.aws.amazon.com/AmazonS3/latest/developerguide/sig-v4-header-based-auth.html + headers: Final = { + "date": "Fri, 24 May 2013 00:00:00 GMT", "host": "examplebucket.s3.amazonaws.com", + "x-amz-content-sha256": "44ce7dd67c959e0d3524ffac1771dfbba87d2b6b4b4e99e42034a8b803f8b072", + "x-amz-date": "20130524T000000Z", "x-amz-storage-class": "REDUCED_REDUNDANCY", + } + signed: Final = "date;host;x-amz-content-sha256;x-amz-date;x-amz-storage-class" + expected: Final = ( + "9e0e90d9c76de8fa5b200d8c849cd5b8dc7a3be3951ddb7f6a76b4158342019d", + "98ad721746da40c64f1a55b78f14c238d841ea1380cd77a1b5971af0ece108bd", + ) + actual: Final = signature("PUT", "/test%24file.text", headers, signed, b"Welcome to Amazon S3.", SECRET, "20130524/us-east-1/s3/aws4_request") + assert actual == expected + assert signature("PUT", "/test$file.text", headers, signed, b"Welcome to Amazon S3.", SECRET, "20130524/us-east-1/s3/aws4_request") != expected + assert encoded_path("/bucket/a=b+c/d e/雪.json") == "/bucket/a%3Db%2Bc/d%20e/%E9%9B%AA.json" + + +@pytest.mark.covers("other.provider_wire.s3.sync_async_reserved_keys_are_signed_and_accepted") +async def test_s3_sync_and_async_uploads_pass_independent_wire_verification(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.integrations.s3_v2 import S3Logger + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement + + monkeypatch.setattr("botocore.auth.get_current_datetime", lambda: datetime(2026, 9, 14)) + payload: Final = {"id": "synthetic-event", "content": "synthetic snow 雪"} + expected_path = "" + + def verify(request: Request) -> Reply: + if request.method != "PUT" or request.target != expected_path: + return Reply(status=403) + try: + authorization: Final = request.headers.get("authorization", "") + assert authorization.startswith("AWS4-HMAC-SHA256 ") + fields: Final = dict(part.split("=", 1) for part in authorization.removeprefix("AWS4-HMAC-SHA256 ").split(", ")) + access, scope = fields["Credential"].split("/", 1) + assert access == ACCESS and scope == "20260914/us-east-1/s3/aws4_request" + assert request.headers["x-amz-date"] == "20260914T000000Z" + signed: Final = fields["SignedHeaders"].split(";") + assert signed == sorted(set(signed)) + assert {"host", "content-md5", "x-amz-date"}.issubset(signed) + assert {name for name in request.headers if name.startswith("x-amz-") and name != "x-amz-content-sha256"}.issubset(signed) + assert request.headers["content-md5"] == base64.b64encode(hashlib.md5(request.body, usedforsecurity=False).digest()).decode() + assert request.headers["x-amz-content-sha256"] == hashlib.sha256(request.body).hexdigest() + expected: Final = signature("PUT", request.target, request.headers, fields["SignedHeaders"], request.body, SECRET, scope)[1] + return Reply(status=200 if hmac.compare_digest(expected, fields["Signature"]) else 403) + except (AssertionError, KeyError, ValueError): + return Reply(status=403) + + with wire_server(verify) as wire: + prior: Final = asyncio.all_tasks() + logger: Final = S3Logger(s3_bucket_name="integration-bucket", s3_region_name="us-east-1", s3_endpoint_url=wire.url, + s3_aws_access_key_id=ACCESS, s3_aws_secret_access_key=SECRET, s3_callback_params_override={}) + owned: Final = asyncio.all_tasks() - prior + assert len(owned) == 1 + try: + for mode in ("sync", "async"): + for key in ("plain.json", "a=b+c/d e/雪.json", "percent%2Fplus+.json"): + expected_path = encoded_path(f"/integration-bucket/{key}") + element: Final = s3BatchLoggingElement(payload=payload, s3_object_key=key, s3_object_download_filename="event.json") + if mode == "sync": + await asyncio.to_thread(logger.upload_data_to_s3, element) + else: + await logger.async_upload_data_to_s3(element) + requests: Final = wire.drain() + assert len(requests) == 1, "Upload must be accepted on its first actual PUT" + request: Final = requests[0] + assert request.target == expected_path + assert json.loads(request.body) == payload + assert verify(request).status == 200 + with httpx.Client(timeout=5, trust_env=False) as client: + corrupt: Final = {**request.headers, "authorization": request.headers["authorization"][:-1] + ("0" if request.headers["authorization"][-1] != "0" else "1")} + assert client.put(wire.url + expected_path, content=request.body, headers=corrupt).status_code == 403 + assert client.put(wire.url + expected_path + "-wrong", content=request.body, headers=request.headers).status_code == 403 + assert client.put(wire.url + expected_path, content=request.body + b" ", headers={name: value for name, value in request.headers.items() if name != "content-length"}).status_code == 403 + fields: Final = dict(part.split("=", 1) for part in request.headers["authorization"].removeprefix("AWS4-HMAC-SHA256 ").split(", ")) + for signed, scope, md5 in ( + (fields["SignedHeaders"].replace("host;", ""), "20260914/us-east-1/s3/aws4_request", request.headers["content-md5"]), + (fields["SignedHeaders"], "20260914/us-west-2/s3/aws4_request", request.headers["content-md5"]), + (fields["SignedHeaders"], "20260914/us-east-1/s3/aws4_request", "AAAAAAAAAAAAAAAAAAAAAA=="), + ): + candidate_headers: Final = {**request.headers, "content-md5": md5} + digest: Final = signature("PUT", request.target, candidate_headers, signed, request.body, SECRET, scope)[1] + candidate_headers["authorization"] = f"AWS4-HMAC-SHA256 Credential={ACCESS}/{scope}, SignedHeaders={signed}, Signature={digest}" + assert client.put(wire.url + expected_path, content=request.body, headers=candidate_headers).status_code == 403 + assert len(wire.drain()) == 6 + + finally: + for task in owned: + task.cancel() + await asyncio.gather(*owned, return_exceptions=True) + assert all(task.done() for task in owned) diff --git a/tests/integration/routing/test_observed_routing.py b/tests/integration/routing/test_observed_routing.py new file mode 100644 index 00000000000..d7398b05fd4 --- /dev/null +++ b/tests/integration/routing/test_observed_routing.py @@ -0,0 +1,98 @@ +import json +import uuid +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml + +from integration._support.client import Gateway, object_value +from integration._support.wire import Reply, Request, wire_server + + +@pytest.mark.covers("other.routing.retries.several_attempts_reach_success_without_hidden_retries", "other.routing.errors.nonretryable_and_exhausted_failures_remain_errors") +def test_retry_counts_and_public_errors_match_actual_provider_attempts(gateway: Gateway) -> None: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, gateway.scenario() as scenario: + original: Final = object_value(gateway.get("/router/settings")["current_values"])["num_retries"] + provider_model: Final = "errors-" + uuid.uuid4().hex + model: Final = scenario.model(model=f"openai/{provider_model}", input_cost_per_token=0, output_cost_per_token=0) + + def remove() -> None: + response: Final = upstream.delete(f"/__scripts/{provider_model}") + assert response.status_code in (200, 404) + assert upstream.get(f"/__scripts/{provider_model}").status_code == 404 + + scenario.cleanups.callback(remove) + try: + for index, (retries, statuses, status, attempts) in enumerate(((2, [500, 500, 200], 200, 3), (2, [400, 200], 400, 1), (1, [429, 429, 200], 429, 2), (1, [500, 500, 200], 500, 2))): + gateway.post("/config/update", {"router_settings": {"num_retries": retries}}) + assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == retries + upstream.post(f"/__scripts/{provider_model}", json={"statuses": statuses}).raise_for_status() + upstream.get("/__observations").raise_for_status() + response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"{provider_model} {index}"}]}) + assert response.status_code == status, response.text + requests: Final = upstream.get("/__observations").json()["requests"] + assert len(requests) == attempts + assert all(request["body"]["model"] == provider_model for request in requests) + assert upstream.get(f"/__scripts/{provider_model}").json()["remaining"] == statuses[attempts:] + if status == 200: + assert response.json()["usage"]["total_tokens"] == 40 + else: + error: Final = response.json()["error"] + assert isinstance(error["message"], str) and "Controlled provider failure" in error["message"] + assert str(error["code"]) == str(status) + assert error["type"] == {400: "invalid_request_error", 429: "throttling_error", 500: "internal_server_error"}[status] + assert error["param"] is None + assert "Traceback" not in response.text and "File \"" not in response.text + finally: + gateway.post("/config/update", {"router_settings": {"num_retries": original}}) + assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == original + + +@pytest.mark.covers("other.routing.fallback.loaded_configuration_selects_only_permitted_target") +def test_loaded_fallback_selects_expected_deployment_and_keeps_response_identity(tmp_path: Path) -> None: + from litellm import Router + + def respond(request: Request) -> Reply: + model: Final = json.loads(request.body)["model"] + assert model in {"primary-wire", "fallback-wire", "unrelated-wire"} + if model == "primary-wire": + return Reply(status=500, body=b'{"error":{"message":"synthetic primary unavailable","type":"api_error","code":"500"}}') + return Reply(body=json.dumps({"id": "response-" + model, "object": "chat.completion", "created": 1, "model": model, "choices": [{"index": 0, "message": {"role": "assistant", "content": "served " + model}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}}).encode()) + + with wire_server(respond) as wire: + path: Final = tmp_path / "fallback.yaml" + path.write_text(yaml.safe_dump({"model_list": [{"model_name": alias, "litellm_params": {"model": "openai/" + upstream, "api_key": "synthetic-routing-key", "api_base": wire.url + "/v1"}} for alias, upstream in (("primary", "primary-wire"), ("fallback", "fallback-wire"), ("unrelated", "unrelated-wire"))], "router_settings": {"num_retries": 0, "disable_cooldowns": True, "fallbacks": [{"primary": ["fallback"]}]}})) + loaded: Final = yaml.safe_load(path.read_text()) + router: Final = Router(model_list=loaded["model_list"], **loaded["router_settings"]) + try: + result: Final = router.completion(model="primary", messages=[{"role": "user", "content": "fallback control"}]) + assert result.id == "response-fallback-wire" + assert result.choices[0].message.content == "served fallback-wire" + assert result.choices[0].finish_reason == "stop" + assert result.usage.prompt_tokens == 11 and result.usage.completion_tokens == 4 + assert tuple(json.loads(request.body)["model"] for request in wire.drain()) == ("primary-wire", "fallback-wire") + control: Final = router.completion(model="unrelated", messages=[{"role": "user", "content": "independent route"}]) + assert control.id == "response-unrelated-wire" + assert tuple(json.loads(request.body)["model"] for request in wire.drain()) == ("unrelated-wire",) + finally: + router.reset() + + +@pytest.mark.covers("other.routing.alias_update.persisted_target_changes_only_selected_route") +def test_saved_deployment_target_update_changes_wire_and_preserves_control(gateway: Gateway) -> None: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, gateway.scenario() as scenario: + prefix: Final = "target-" + uuid.uuid4().hex + model: Final = scenario.model(model="openai/" + prefix + "-first", input_cost_per_token=0, output_cost_per_token=0) + other: Final = scenario.model(model="openai/" + prefix + "-control", input_cost_per_token=0, output_cost_per_token=0) + target: Final = next(entry for entry in gateway.get("/model/info")["data"] if entry["model_name"] == model) + for generation, suffix in enumerate(("first", "second")): + if generation: + response: Final = gateway.request("PATCH", f"/model/{target['model_info']['id']}/update", {"litellm_params": {"model": "openai/" + prefix + "-second"}}) + assert response.status_code == 200, response.text + upstream.get("/__observations").raise_for_status() + for alias in (model, other): + assert gateway.chat(alias, text=f"{prefix} generation {generation}")["usage"]["total_tokens"] == 40 + requests: Final = upstream.get("/__observations").json()["requests"] + assert [request["body"]["model"] for request in requests] == [prefix + "-" + suffix, prefix + "-control"] diff --git a/tests/integration/routing/test_redis_recovery.py b/tests/integration/routing/test_redis_recovery.py new file mode 100644 index 00000000000..81d27a190b0 --- /dev/null +++ b/tests/integration/routing/test_redis_recovery.py @@ -0,0 +1,59 @@ +import os +import uuid +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +import httpx +import psycopg +import pytest +from psycopg import sql +from redis import Redis + +from integration._support.client import Gateway, eventually +from integration._support.process import owned_proxy +from integration._support.redis_process import owned_redis + + +@pytest.mark.covers("other.routing.redis.owned_outage_recovers_serving_and_response_cache") +def test_owned_redis_outage_recovers_requests_and_real_response_cache(gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + original: Final = os.environ["DATABASE_URL"] + identity: Final = "integration_recovery_" + uuid.uuid4().hex + parsed: Final = urlsplit(original) + database_url: Final = urlunsplit((parsed.scheme, parsed.netloc, "/" + identity, "", "")) + with psycopg.connect(original, autocommit=True) as admin: + admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(identity))) + try: + with owned_redis(tmp_path) as cache, monkeypatch.context() as environment: + environment.setenv("DATABASE_URL", database_url) + with owned_proxy(gateway, tmp_path, {"DATABASE_URL": database_url, "REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "1"}) as candidate, candidate.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + for generation in ("before", "after"): + with Redis(host=cache.host, port=cache.port, socket_timeout=1) as client: + eventually(client.ping, bool) + eventually(lambda: client.pubsub_numsub("litellm_proxy.auth_cache_invalidation")[0][1], lambda count: count >= 1, seconds=8) + upstream.get("/__observations").raise_for_status() + first: Final = candidate.chat(model, key=key, text=identity + generation) + second: Final = candidate.chat(model, key=key, text=identity + generation) + assert first["id"] == second["id"] + assert first["choices"] == second["choices"] and first["usage"]["total_tokens"] == 40 + assert len(upstream.get("/__observations").json()["requests"]) == 1 + with Redis(host=cache.host, port=cache.port, socket_timeout=1) as client: + eventually( + lambda first=first: tuple(client.get(name) for name in client.scan_iter() if client.type(name) == b"string"), + lambda values, first=first: any(str(first["id"]).encode() in value for value in values if value is not None), + seconds=10, + ) + if generation == "before": + cache.stop() + upstream.get("/__observations").raise_for_status() + during: Final = candidate.chat(model, key=key, text=identity + "during") + assert during["usage"]["total_tokens"] == 40 + assert len(upstream.get("/__observations").json()["requests"]) == 1 + cache.start() + with psycopg.connect(database_url) as fresh: + assert fresh.execute('SELECT count(*) FROM "LiteLLM_VerificationToken"').fetchone()[0] >= 1 + finally: + admin.execute(sql.SQL("DROP DATABASE {}").format(sql.Identifier(identity))) + assert admin.execute("SELECT datname FROM pg_database WHERE datname=%s", (identity,)).fetchall() == [] diff --git a/tests/integration/streaming/test_stream_contracts.py b/tests/integration/streaming/test_stream_contracts.py new file mode 100644 index 00000000000..0c0fd8bc47c --- /dev/null +++ b/tests/integration/streaming/test_stream_contracts.py @@ -0,0 +1,149 @@ +import asyncio +import json +import threading +import uuid +from typing import Final + +import pytest +from hypothesis import Phase, example, given, settings, strategies as st +from openai import OpenAI + +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, wire_server + + +def frame(identity: str, delta: dict, *, finish: str | None = None) -> bytes: + value: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]} + return b"data: " + json.dumps(value, ensure_ascii=False).encode() + b"\n\n" + + +def text_stream(identity: str) -> tuple[bytes, ...]: + usage: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [], "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}} + return (frame(identity, {"role": "assistant", "content": "Hello "}), frame(identity, {"content": "雪 café"}), frame(identity, {}, finish="stop"), b"data: " + json.dumps(usage).encode() + b"\n\n", b"data: [DONE]\n\n") + + +@pytest.mark.covers("other.streaming.byte_partitions.preserve_text_identity_and_usage") +def test_generated_tcp_partitions_preserve_unicode_text_identity_and_final_usage() -> None: + import litellm + + body: Final = b"".join(text_stream("stream-partition-control")) + + @settings(max_examples=20, deadline=None, database=None, phases=(Phase.explicit, Phase.generate, Phase.shrink)) + @example(cuts=tuple(range(1, len(body)))) + @example(cuts=()) + @given(cuts=st.lists(st.integers(min_value=1, max_value=len(body) - 1), max_size=35, unique=True).map(tuple)) + def check(cuts: tuple[int, ...]) -> None: + boundaries: Final = (0, *sorted(cuts), len(body)) + pieces: Final = tuple(body[left:right] for left, right in zip(boundaries, boundaries[1:])) + with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=pieces)) as wire: + stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "partition control"}], stream=True, stream_options={"include_usage": True}, timeout=5, num_retries=0) + try: + chunks: Final = tuple(stream) + finally: + asyncio.run(stream.aclose()) + assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café" + assert {chunk.id for chunk in chunks} == {"stream-partition-control"} + assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == ["stop"] + usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None) + assert len(usages) == 1 + assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4 + assert len(wire.drain()) == 1 + + check() + + +@pytest.mark.covers("other.streaming.tools.fragmented_calls_keep_independent_arguments") +def test_fragmented_tool_names_and_arguments_keep_each_call_identity() -> None: + import litellm + + identity: Final = "stream-tools-control" + deltas: Final = ( + {"role": "assistant", "tool_calls": [{"index": 0, "id": "call-add", "type": "function", "function": {"name": "ad", "arguments": ""}}, {"index": 1, "id": "call-multiply", "type": "function", "function": {"name": "multi", "arguments": ""}}]}, + {"tool_calls": [{"index": 1, "function": {"name": "ply", "arguments": '{"x":3,'}}, {"index": 0, "function": {"arguments": '{"x":1,'}}]}, + {"tool_calls": [{"index": 0, "function": {"name": "d", "arguments": '"y":2}'}}, {"index": 1, "function": {"arguments": '"y":4}'}}]}, + ) + frames: Final = (*tuple(frame(identity, delta) for delta in deltas), frame(identity, {}, finish="tool_calls"), b"data: [DONE]\n\n") + with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire: + stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "tool control"}], stream=True, timeout=5, num_retries=0) + try: + chunks: Final = tuple(stream) + finally: + asyncio.run(stream.aclose()) + events: Final = tuple((choice.index, tool) for chunk in chunks for choice in chunk.choices for tool in (choice.delta.tool_calls or ())) + for index, name, call_id, arguments in ((0, "add", "call-add", {"x": 1, "y": 2}), (1, "multiply", "call-multiply", {"x": 3, "y": 4})): + selected: Final = tuple(tool for choice, tool in events if (choice, tool.index) == (0, index)) + assert "".join(tool.id or "" for tool in selected) == call_id + assert "".join(tool.function.name or "" for tool in selected) == name + assert json.loads("".join(tool.function.arguments or "" for tool in selected)) == arguments + assert {tool.index for _, tool in events} == {0, 1} + assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == ["tool_calls"] + assert len(wire.drain()) == 1 + + +@pytest.mark.covers("other.streaming.usage.client_visibility_preserves_persisted_accounting") +def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + for include in (None, False, True): + identity: Final = "stream-usage-" + uuid.uuid4().hex + with wire_server(lambda request, identity=identity: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire: + model: Final = scenario.model(api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002) + with OpenAI(api_key=gateway.key, base_url=str(gateway.client.base_url), timeout=5, max_retries=0) as client: + stream: Final = client.chat.completions.create(model=model, messages=[{"role": "user", "content": identity}], stream=True, **({} if include is None else {"stream_options": {"include_usage": include}})) + with stream: + chunks: Final = tuple(stream) + assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café" + assert {chunk.id for chunk in chunks} == {identity} + usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None) + assert len(usages) == (1 if include else 0) + if include: + assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4 + requests: Final = wire.drain() + assert len(requests) == 1 + assert json.loads(requests[0].body)["stream_options"]["include_usage"] is True + rows: Final = eventually(lambda identity=identity: read_rows('SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), lambda values: len(values) == 1, seconds=70) + assert rows[0]["prompt_tokens"] == 11 and rows[0]["completion_tokens"] == 4 + assert float(rows[0]["spend"]) == pytest.approx(0.019) + + +@pytest.mark.covers("other.streaming.failure.truncated_transport_raises_and_control_recovers") +def test_truncated_http_stream_is_an_error_and_next_stream_succeeds() -> None: + import litellm + + for truncated in (True, False): + with wire_server(lambda request, truncated=truncated: Reply(content_type="text/event-stream", chunks=text_stream("stream-truncated"), abort_after=1 if truncated else None)) as wire: + stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "truncation control"}], stream=True, timeout=5, num_retries=0) + try: + if truncated: + with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="incomplete chunked read") as failure: + tuple(stream) + assert isinstance(failure.value.original_exception, litellm.APIConnectionError) + assert failure.value.generated_content == "Hello " + assert failure.value.is_pre_first_chunk is False + else: + chunks: Final = tuple(stream) + assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café" + assert any(choice.finish_reason == "stop" for chunk in chunks for choice in chunk.choices) + finally: + asyncio.run(stream.aclose()) + assert len(wire.drain()) == 1 + + +@pytest.mark.covers("other.streaming.cancellation.closes_actual_provider_connection") +def test_client_cancellation_releases_the_actual_provider_connection() -> None: + import litellm + + gate: Final = threading.Event() + frames: Final = (frame("stream-cancel", {"role": "assistant", "content": "first"}), b":" + b"x" * 4_000_000 + b"\n\n", b"data: [DONE]\n\n") + with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames, gate_after_first=gate)) as wire: + stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "cancellation control"}], stream=True, timeout=5, num_retries=0) + try: + first: Final = next(stream) + assert first.choices[0].delta.content == "first" + finally: + try: + asyncio.run(stream.aclose()) + finally: + gate.set() + assert wire.disconnected.get(timeout=5) == "/v1/chat/completions" + assert len(wire.drain()) == 1 From ba03f60f7143f5d5b78228b63c1a0fe8043642f8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 19:50:28 +0000 Subject: [PATCH 116/187] fix(models): add batch and flex tier prices on Azure dated and Gemini latest alias rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 16 ++++++++++++++-- model_prices_and_context_window.json | 16 ++++++++++++++-- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2575b227bfc..0f3535d9b18 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -8806,6 +8806,7 @@ "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, "cache_read_input_token_cost_flex": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, "input_cost_per_token_above_272k_tokens_flex": 5e-06, "input_cost_per_token_priority": 1.25e-05, @@ -8817,6 +8818,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, + "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, @@ -8907,9 +8909,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, "input_cost_per_token": 5.5e-06, + "input_cost_per_token_batches": 2.75e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, "output_cost_per_token": 3.3e-05, + "output_cost_per_token_batches": 1.65e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, "output_cost_per_token_priority": 8.25e-05, "litellm_provider": "azure", @@ -9004,9 +9008,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, "input_cost_per_token": 5.5e-06, + "input_cost_per_token_batches": 2.75e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, "output_cost_per_token": 3.3e-05, + "output_cost_per_token_batches": 1.65e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, "output_cost_per_token_priority": 8.25e-05, "litellm_provider": "azure", @@ -56989,7 +56995,10 @@ "prompt_cache_min_tokens": 4096, "supports_native_streaming": true, "supports_url_context": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_read_input_token_cost_flex": 2e-07, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 6e-06 }, "gemini/gemini-pro-latest": { "cache_read_input_token_cost": 2e-07, @@ -57048,7 +57057,10 @@ "prompt_cache_min_tokens": 4096, "supports_native_streaming": true, "supports_url_context": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_read_input_token_cost_flex": 2e-07, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 6e-06 }, "gemini-exp-1206": { "cache_read_input_token_cost": 3e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2575b227bfc..0f3535d9b18 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -8806,6 +8806,7 @@ "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, "cache_read_input_token_cost_flex": 2.5e-07, "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, "input_cost_per_token_above_272k_tokens_flex": 5e-06, "input_cost_per_token_priority": 1.25e-05, @@ -8817,6 +8818,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, + "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens": 4.5e-05, "output_cost_per_token_priority": 7.5e-05, "output_cost_per_token_above_272k_tokens_priority": 9e-05, @@ -8907,9 +8909,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, "input_cost_per_token": 5.5e-06, + "input_cost_per_token_batches": 2.75e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, "output_cost_per_token": 3.3e-05, + "output_cost_per_token_batches": 1.65e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, "output_cost_per_token_priority": 8.25e-05, "litellm_provider": "azure", @@ -9004,9 +9008,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, "input_cost_per_token": 5.5e-06, + "input_cost_per_token_batches": 2.75e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, "output_cost_per_token": 3.3e-05, + "output_cost_per_token_batches": 1.65e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, "output_cost_per_token_priority": 8.25e-05, "litellm_provider": "azure", @@ -56989,7 +56995,10 @@ "prompt_cache_min_tokens": 4096, "supports_native_streaming": true, "supports_url_context": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_read_input_token_cost_flex": 2e-07, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 6e-06 }, "gemini/gemini-pro-latest": { "cache_read_input_token_cost": 2e-07, @@ -57048,7 +57057,10 @@ "prompt_cache_min_tokens": 4096, "supports_native_streaming": true, "supports_url_context": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "cache_read_input_token_cost_flex": 2e-07, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 6e-06 }, "gemini-exp-1206": { "cache_read_input_token_cost": 3e-08, From 9ebd55e53e4e1a9347bf593d99915c30e0d3709c Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 19:57:55 +0000 Subject: [PATCH 117/187] test(router): cover success callback recovering the count when the pre-header increment fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_router_helper_utils.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 4a02ec62064..3690143d4ad 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1127,6 +1127,45 @@ async def test_success_callback_running_during_pre_header_increment_does_not_dou assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1) +class _FailingFirstIncrementCache(DualCache): + def __init__(self) -> None: + super().__init__(in_memory_cache=InMemoryCache()) + self.increment_calls = 0 + + async def async_increment_cache_pipeline( + self, + increment_list: list[RedisPipelineIncrementOperation], + local_only: bool = False, + parent_otel_span: object = None, + **kwargs: object, + ) -> list[float] | None: + self.increment_calls += 1 + if self.increment_calls == 1: + raise RuntimeError("cache unavailable") + return await super().async_increment_cache_pipeline( + increment_list, local_only=local_only, parent_otel_span=parent_otel_span, **kwargs + ) + + +@pytest.mark.asyncio +async def test_success_callback_counts_fully_when_pre_header_increment_fails(): + router = _rpm_tpm_router("lit-3058-recover") + cache = _FailingFirstIncrementCache() + router.cache = cache + + response = await router.acompletion( + model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong" + ) + + expected = (response.usage.total_tokens, 1) + for _ in range(50): + if await router.get_model_group_usage("gpt-5-mini") == expected: + break + await asyncio.sleep(0.1) + assert await router.get_model_group_usage("gpt-5-mini") == expected + assert cache.increment_calls == 2 + + @pytest.mark.asyncio async def test_increment_deployment_usage_for_response_skips_session_wrappers(): router = _rpm_tpm_router("lit-3058-ws") From a229f99ac8f924b8d3499186a97bc8b5fb3cff29 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 19:59:07 +0000 Subject: [PATCH 118/187] test(models): drive Grok prompt caching coverage through litellm APIs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...tex_ai_xai_grok_prompt_caching_metadata.py | 38 +++++-------------- 1 file changed, 9 insertions(+), 29 deletions(-) diff --git a/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py b/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py index b8b3897f004..72e98711f0c 100644 --- a/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py +++ b/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py @@ -1,37 +1,28 @@ -from pathlib import Path from typing import Final import pytest -from pydantic import TypeAdapter +import litellm from litellm import get_model_info from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.utils import supports_prompt_caching -REPO_ROOT: Final = Path(__file__).parents[2] MODEL: Final = "vertex_ai/xai/grok-4.6" GROK_KEY_PREFIXES: Final = ("vertex_ai/xai/grok-", "azure_ai/grok-", "xai/grok-") -COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) -MAIN_COST_MAP: Final = REPO_ROOT / "model_prices_and_context_window.json" -BACKUP_COST_MAP: Final = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" -def _cost_map(path: Path) -> dict[str, dict[str, object]]: - return COST_MAP_ADAPTER.validate_json(path.read_bytes()) - - -def test_grok_entries_with_cache_read_price_advertise_prompt_caching() -> None: - cost_map = _cost_map(MAIN_COST_MAP) - cached_grok_entries = tuple( - (key, entry) - for key, entry in cost_map.items() +@pytest.mark.usefixtures("local_model_cost_map") +def test_grok_models_with_cache_read_price_advertise_prompt_caching() -> None: + cached_grok_models = tuple( + key + for key, entry in litellm.model_cost.items() if key.startswith(GROK_KEY_PREFIXES) and entry.get("cache_read_input_token_cost") ) - assert cached_grok_entries, "expected at least one grok entry with a cache read price" + assert cached_grok_models, "expected at least one grok model with a cache read price" - missing_flag = tuple(key for key, entry in cached_grok_entries if entry.get("supports_prompt_caching") is not True) + missing_flag = tuple(key for key in cached_grok_models if supports_prompt_caching(model=key) is not True) assert missing_flag == (), ( - f"grok entries with cache_read_input_token_cost missing supports_prompt_caching: {missing_flag}" + f"grok models with cache_read_input_token_cost fail supports_prompt_caching: {missing_flag}" ) @@ -45,14 +36,3 @@ def test_vertex_ai_grok_4_6_supports_prompt_caching_via_get_model_info() -> None assert info.get("supports_prompt_caching") is True assert supports_prompt_caching(model=MODEL) is True - - -def test_vertex_ai_grok_entries_source_and_backup_match() -> None: - main_map = _cost_map(MAIN_COST_MAP) - backup_map = _cost_map(BACKUP_COST_MAP) - - vertex_grok_keys = tuple(key for key in main_map if key.startswith("vertex_ai/xai/grok-")) - assert vertex_grok_keys, "expected at least one vertex_ai/xai/grok- entry" - - mismatched = tuple(key for key in vertex_grok_keys if backup_map.get(key) != main_map[key]) - assert mismatched == (), f"vertex grok entries differ between source and backup: {mismatched}" From 19795339012f9a83a480298f6926b857ea753c81 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 20:10:22 +0000 Subject: [PATCH 119/187] test(router): force the callback to observe the usage stamp before the pre-header increment fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_router_helper_utils.py | 42 ++++++++++++------- 1 file changed, 27 insertions(+), 15 deletions(-) diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 3690143d4ad..bb389693311 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1127,9 +1127,11 @@ async def test_success_callback_running_during_pre_header_increment_does_not_dou assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1) -class _FailingFirstIncrementCache(DualCache): +class _UnavailableIncrementCache(DualCache): def __init__(self) -> None: super().__init__(in_memory_cache=InMemoryCache()) + self.first_increment_started = asyncio.Event() + self.release_first_increment = asyncio.Event() self.increment_calls = 0 async def async_increment_cache_pipeline( @@ -1141,29 +1143,39 @@ class _FailingFirstIncrementCache(DualCache): ) -> list[float] | None: self.increment_calls += 1 if self.increment_calls == 1: - raise RuntimeError("cache unavailable") - return await super().async_increment_cache_pipeline( - increment_list, local_only=local_only, parent_otel_span=parent_otel_span, **kwargs - ) + self.first_increment_started.set() + await self.release_first_increment.wait() + raise RuntimeError("cache unavailable") @pytest.mark.asyncio -async def test_success_callback_counts_fully_when_pre_header_increment_fails(): - router = _rpm_tpm_router("lit-3058-recover") - cache = _FailingFirstIncrementCache() +async def test_callback_observing_stamp_before_pre_header_increment_fails_leaves_no_stamp_behind(): + router = _rpm_tpm_router("lit-3058-fail") + cache = _UnavailableIncrementCache() router.cache = cache + metadata: dict[str, object] = {} - response = await router.acompletion( - model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong" + request = asyncio.ensure_future( + router.acompletion( + model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong", metadata=metadata + ) ) - - expected = (response.usage.total_tokens, 1) + await asyncio.wait_for(cache.first_increment_started.wait(), timeout=5) + assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 30 for _ in range(50): - if await router.get_model_group_usage("gpt-5-mini") == expected: + if get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1: break await asyncio.sleep(0.1) - assert await router.get_model_group_usage("gpt-5-mini") == expected - assert cache.increment_calls == 2 + assert get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1 + assert cache.increment_calls == 1 + + cache.release_first_increment.set() + response = await request + + assert response.usage.total_tokens == 30 + assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in metadata + assert _ratelimit_headers(response)["x-ratelimit-remaining-requests"] == 100 + assert await router.get_model_group_usage("gpt-5-mini") == (None, None) @pytest.mark.asyncio From fa5d31a8378f53cb4010da0b773b5d59e55da0b3 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 13:14:47 -0700 Subject: [PATCH 120/187] test(integration): send OpenAI-shaped error types from the fake upstream Since #40994 the proxy relays the upstream error body on a 400, so the public error type is now whatever the upstream sent instead of the status-derived name. The fake upstream answered every scripted failure with type api_error, which made the public-error contract in test_retry_counts_and_public_errors_match_actual_provider_attempts fail on main. The fake now sends the type a real OpenAI-compatible upstream sends for the status, so the assertion holds whether the proxy relays or maps the type --- tests/integration/_support/upstream.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 8bc4100abfd..04a6ea02eec 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -31,6 +31,12 @@ INTERNAL_FIELDS: Final = frozenset( ) +def error_type(status: int) -> str: + if status == 429: + return "rate_limit_error" + return "invalid_request_error" if status < 500 else "server_error" + + @dataclass(frozen=True, slots=True) class Observation: path: str @@ -66,7 +72,7 @@ class Provider: status: Final = script.popleft() if status != 200: return JSONResponse( - {"error": {"message": "Controlled provider failure", "type": "api_error", "code": str(status)}}, + {"error": {"message": "Controlled provider failure", "type": error_type(status), "code": str(status)}}, status_code=status, ) return await chat_completions(request) From df9a87f44a09d77b10b3868700ae6569bdedb57d Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 20:28:03 +0000 Subject: [PATCH 121/187] fix(otel): keep caller tracestate on the legacy request span Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/opentelemetry.py | 11 +--- .../integrations/test_opentelemetry.py | 59 +++++++++++++++++++ 2 files changed, 61 insertions(+), 9 deletions(-) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index d4e7fcb577e..2623a9b5a56 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -2932,16 +2932,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) propagator: Final = TraceContextTextMapPropagator() - carrier: Final = {"traceparent": _traceparent} + carrier: Final = {key: headers[key] for key in ("traceparent", "tracestate") if headers.get(key) is not None} _parent_context: Final = propagator.extract(carrier=carrier) return _parent_context def _get_span_context(self, kwargs, default_span: Span | None = None): from opentelemetry import context, trace - from opentelemetry.trace.propagation.tracecontext import ( - TraceContextTextMapPropagator, - ) litellm_params: Final = kwargs.get("litellm_params", {}) or {} proxy_server_request: Final = litellm_params.get("proxy_server_request", {}) or {} @@ -2965,11 +2962,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # Priority 2: HTTP traceparent header if traceparent is not None: verbose_logger.debug("OpenTelemetry: Using traceparent header for context propagation") - carrier: Final = {"traceparent": traceparent} - return ( - TraceContextTextMapPropagator().extract(carrier=carrier), - None, - ) + return self.get_traceparent_from_header(headers=headers), None # Priority 3: Active span from global context (auto-detection) try: diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 9ec8489f784..443b9012235 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -5423,6 +5423,65 @@ class TestGetSpanContextLitellmMetadataFallback(unittest.TestCase): self.assertIsNone(detected_span) +class TestInboundTraceContextKeepsCallerTracestate(unittest.TestCase): + """The request span built from inbound W3C headers must carry the caller's + tracestate so outbound propagation (passthrough) re-emits it instead of + dropping it alongside the stripped stale header.""" + + CALLER_TRACEPARENT = "00-" + "a" * 32 + "-" + "b" * 16 + "-01" + CALLER_TRACESTATE = "vendor=abc,other=xyz" + + def _otel(self): + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + otel = OpenTelemetry() + otel.tracer = provider.get_tracer(__name__) + return otel + + def test_request_span_propagates_caller_tracestate_downstream(self): + from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + + from litellm.integrations.otel.plumbing.context import inject_trace_context + + inbound = {"traceparent": self.CALLER_TRACEPARENT, "tracestate": self.CALLER_TRACESTATE} + span = self._otel().create_litellm_proxy_request_started_span( + start_time=datetime.now(timezone.utc), headers=inbound + ) + outbound = inject_trace_context(inbound, parent_span=span) + span.end() + + propagated = trace.get_current_span(TraceContextTextMapPropagator().extract(outbound)).get_span_context() + self.assertEqual(outbound["tracestate"], self.CALLER_TRACESTATE) + self.assertEqual(propagated.trace_id, span.get_span_context().trace_id) + self.assertEqual(propagated.span_id, span.get_span_context().span_id) + self.assertNotEqual(outbound["traceparent"], self.CALLER_TRACEPARENT) + + def test_request_span_without_caller_tracestate_emits_none(self): + from litellm.integrations.otel.plumbing.context import inject_trace_context + + inbound = {"traceparent": self.CALLER_TRACEPARENT} + span = self._otel().create_litellm_proxy_request_started_span( + start_time=datetime.now(timezone.utc), headers=inbound + ) + outbound = inject_trace_context(inbound, parent_span=span) + span.end() + + self.assertNotIn("tracestate", outbound) + self.assertNotEqual(outbound["traceparent"], self.CALLER_TRACEPARENT) + + def test_span_context_from_header_keeps_caller_tracestate(self): + kwargs = { + "litellm_params": { + "proxy_server_request": { + "headers": {"traceparent": self.CALLER_TRACEPARENT, "tracestate": self.CALLER_TRACESTATE} + } + } + } + ctx, detected_span = self._otel()._get_span_context(kwargs) + self.assertIsNone(detected_span) + self.assertEqual(trace.get_current_span(ctx).get_span_context().trace_state.to_header(), self.CALLER_TRACESTATE) + + class TestEndProxySpanLitellmMetadataFallback(unittest.TestCase): """ Tests for _end_proxy_span_from_kwargs() falling back to litellm_metadata. From 8acd2477a65cf6c36346169afd7b0c644833ab1f Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 20:34:20 +0000 Subject: [PATCH 122/187] fix(guardrails): hold legacy function_call stream windows until the end-of-stream scan Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/guardrail_translation/handler.py | 8 +- .../test_openai_guardrail_handler.py | 22 ++++ .../test_streaming_buffer_until_moderated.py | 100 +++++++++++++----- 3 files changed, 104 insertions(+), 26 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 9b818696bbb..5b232dd0bf4 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -806,7 +806,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): stream_item_fingerprint(tool_call) for chunk in responses_so_far for choice in _stream_chunk_choices(chunk) - for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls") + for tool_call in _streamed_delta_tool_calls(stream_item_field(choice, "delta")) ) @staticmethod @@ -1344,6 +1344,12 @@ def _stream_chunk_choices(item: object) -> Sequence[object]: return () +def _streamed_delta_tool_calls(delta: object) -> tuple[object, ...]: + function_call: Final = stream_item_field(delta, "function_call") + legacy: Final = () if function_call is None else (function_call,) + return stream_item_items(delta, "tool_calls") + legacy + + def _blocked_stream_identity( exc: "ModifyResponseException", responses_so_far: Sequence[object] ) -> tuple[str, int, str]: 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 0e59e1d1590..9b88161aca9 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 @@ -2213,6 +2213,28 @@ class TestStreamingScanKey: assert ended_key.tool_calls_in_flight is False assert ended_key != open_key + def test_legacy_function_call_delta_is_held_like_a_tool_call(self): + from litellm.types.utils import Delta, FunctionCall, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + function_chunk = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=None, function_call=FunctionCall(name="run_shell", arguments='{"cmd": "rm"}')), + finish_reason=None, + ) + ] + ) + open_key = handler.get_streaming_scan_key([self._chunk("hi"), function_chunk]) + ended_key = handler.get_streaming_scan_key( + [self._chunk("hi"), function_chunk, self._chunk(None, finish_reason="function_call")] + ) + assert open_key.tool_calls_in_flight is True + assert open_key.tool_calls == () + assert len(ended_key.tool_calls) == 1 and "run_shell" in ended_key.tool_calls[0] + assert ended_key.tool_calls_in_flight is False + def test_text_after_the_first_choice_finishes_still_changes_the_key(self): handler = OpenAIChatCompletionsHandler() first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py index 318488c6d7b..db937f18e96 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py @@ -29,6 +29,7 @@ from litellm.types.utils import ( ChatCompletionDeltaToolCall, Delta, Function, + FunctionCall, GenericGuardrailAPIInputs, ModelResponseStream, StreamingChoices, @@ -219,29 +220,30 @@ def _chat_chunk(content: str = "", finish_reason: str | None = None) -> ModelRes ) -def _tool_call_chunk(arguments: str, finish_reason: str | None = None) -> ModelResponseStream: +def _tool_call_chunk( + arguments: str, finish_reason: str | None = None, legacy_function_call: bool = False +) -> ModelResponseStream: + delta = ( + Delta(role="assistant", content=None, function_call=FunctionCall(name="run_shell", arguments=arguments)) + if legacy_function_call + else Delta( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_1", + type="function", + index=0, + function=Function(name="run_shell", arguments=arguments), + ) + ], + ) + ) return ModelResponseStream( id="chatcmpl-windowed", created=1724900000, model="gpt-4", - choices=[ - StreamingChoices( - index=0, - delta=Delta( - role="assistant", - content=None, - tool_calls=[ - ChatCompletionDeltaToolCall( - id="call_1", - type="function", - index=0, - function=Function(name="run_shell", arguments=arguments), - ) - ], - ), - finish_reason=finish_reason, - ) - ], + choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)], ) @@ -250,13 +252,14 @@ async def _windowed_chat_stream( collected: List[Any], content_chunks: List[str], tool_argument_chunks: List[str] | None = None, + legacy_function_call: bool = False, ) -> AsyncGenerator[ModelResponseStream, None]: for content in content_chunks: yielded_count.append(len(collected)) yield _chat_chunk(content) for arguments in tool_argument_chunks or []: yielded_count.append(len(collected)) - yield _tool_call_chunk(arguments) + yield _tool_call_chunk(arguments, legacy_function_call=legacy_function_call) yielded_count.append(len(collected)) yield _chat_chunk(finish_reason="tool_calls" if tool_argument_chunks else "stop") @@ -271,11 +274,22 @@ def _tool_argument_text(chunks: List[Any]) -> str: ) +def _function_call_argument_text(chunks: list[Any]) -> str: + return "".join( + choice.delta.function_call.arguments or "" + for chunk in chunks + if isinstance(chunk, ModelResponseStream) + for choice in chunk.choices + if choice.delta.function_call is not None + ) + + async def _run_windowed( guardrail: CustomGuardrail, content_chunks: List[str], end_of_stream_only: bool = False, tool_argument_chunks: List[str] | None = None, + legacy_function_call: bool = False, ) -> tuple[List[Any], List[int]]: guardrail.streaming_buffer_until_moderated = True guardrail.streaming_buffer_release_on_scan = True @@ -292,7 +306,9 @@ async def _run_windowed( yielded_count: List[int] = [] async for chunk in unified.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, - response=_windowed_chat_stream(yielded_count, collected, content_chunks, tool_argument_chunks), + response=_windowed_chat_stream( + yielded_count, collected, content_chunks, tool_argument_chunks, legacy_function_call + ), request_data=request_data, ): collected.append(chunk) @@ -305,13 +321,24 @@ def _responses_message_stream_events(text_chunks: List[str]) -> List[dict]: return [ {"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}}, *( - {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": text} + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } for text in text_chunks ), {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": content}}, { "type": "response.completed", - "response": {"id": "resp_1", "model": "gpt-4o", "status": "completed", "output": [{**message, "content": content}]}, + "response": { + "id": "resp_1", + "model": "gpt-4o", + "status": "completed", + "output": [{**message, "content": content}], + }, }, ] @@ -322,7 +349,13 @@ def _responses_truncated_function_call_events(text: str, argument_chunks: List[s function_call = {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "run_shell"} return [ {"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}}, - {"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": text}, + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + }, {"type": "response.output_item.added", "output_index": 1, "item": {**function_call, "arguments": ""}}, *( {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": arguments} @@ -463,6 +496,21 @@ async def test_windowed_buffer_holds_tool_call_windows_until_end_of_stream_scan( assert guardrail.tool_call_scan_indexes == [guardrail.scan_count] +@pytest.mark.asyncio +async def test_windowed_buffer_holds_legacy_function_call_windows_until_end_of_stream(): + guardrail = _PassingGuardrail(guardrail_name="windowed-functions", event_hook="post_call") + content_chunks = ["one ", "two ", "three "] + function_argument_chunks = ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'] + + collected, yielded_count = await _run_windowed( + guardrail, content_chunks, tool_argument_chunks=function_argument_chunks, legacy_function_call=True + ) + + assert yielded_count == [0, 0, 2, 2, 2, 2, 2] + assert _chat_text(collected) == "".join(content_chunks) + assert _function_call_argument_text(collected) == "".join(function_argument_chunks) + + def test_tool_call_only_scan_key_is_not_skipped_as_empty(): assert _is_redundant_scan(StreamingScanKey(texts=("",)), None) is True assert _is_redundant_scan(StreamingScanKey(texts=("",), tool_calls=("run_shell:{}",)), None) is False @@ -470,7 +518,9 @@ def test_tool_call_only_scan_key_is_not_skipped_as_empty(): @pytest.mark.asyncio async def test_windowed_responses_output_item_done_round_keeps_text_window_withheld(): - guardrail = _MarkerBlockingGuardrail(guardrail_name="windowed-responses", event_hook="post_call", marker=ORIGINAL_MARKER) + guardrail = _MarkerBlockingGuardrail( + guardrail_name="windowed-responses", event_hook="post_call", marker=ORIGINAL_MARKER + ) events = _responses_message_stream_events(["one ", f"{ORIGINAL_MARKER} "]) raw = await _run_windowed_responses(guardrail, events) From 9421b26bf6dcc98bee10e2cbbd45bf6ff1c23166 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 13:47:07 -0700 Subject: [PATCH 123/187] fix(e2e): stage the seeded device id per thread, not per process Build 232 took two compat cells red with a FileNotFoundError renaming `.claude.json.197` onto `.claude.json`. `run_claude_models_parallel` drives several models from one process, so a pid-suffixed staged name is shared between threads: one thread renamed the file the other was still writing, and the loser died on a path that no longer existed. mkstemp in the same directory gives a name that is unique per thread as well as per process, and the rename stays atomic. --- .../test_request_determinism.py | 19 ++++++++++++++++++- tests/e2e/claude_code/cli_driver.py | 11 ++++++++--- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py index 09a181162d2..5046f35c73b 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py @@ -33,7 +33,7 @@ from typing import List, Tuple import pytest -from claude_code.cli_driver import _stable_cli_state, run_claude +from claude_code.cli_driver import _FIXED_CLI_USER_ID, _seed_cli_identity, _stable_cli_state, run_claude from claude_code.rate_limiter import RateLimiter _STUB_REPLY = { @@ -141,3 +141,20 @@ def test_concurrent_cells_do_not_collide_on_the_pinned_session( assert codes == [0, 0, 0, 0] assert bodies, "the CLI sent no request to the stub, so there is nothing to compare" assert set(Counter(bodies).values()) == {4} + + +def test_seeding_the_device_id_survives_threads_racing_on_the_same_directory(tmp_path: Path) -> None: + """`run_claude_models_parallel` drives several models from one process, so the + seed's staged file has to be unique per thread and not merely per process.""" + config_dir = tmp_path / "config" + config_dir.mkdir() + seeded = config_dir / ".claude.json" + + for _round in range(20): + seeded.unlink(missing_ok=True) + with ThreadPoolExecutor(max_workers=16) as pool: + for outcome in [pool.submit(_seed_cli_identity, str(config_dir)) for _ in range(16)]: + outcome.result() + + assert json.loads(seeded.read_text(encoding="utf-8"))["userID"] == _FIXED_CLI_USER_ID + assert sorted(entry.name for entry in config_dir.iterdir()) == [".claude.json"] diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index 3fad87c7479..a01d8ab3e7c 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -143,7 +143,12 @@ def _seed_cli_identity(config_dir: str) -> None: `userID`, and sends them in `metadata.user_id` forever after, so the value is stable for exactly as long as that file lives. Pinning it, and the session id passed beside it, costs nothing: both feed abuse detection - rather than quota, caching or continuity.""" + rather than quota, caching or continuity. + + The staged name has to be unique per *thread*, not per process: + `run_claude_models_parallel` drives several models from one process, so a + pid-suffixed name lets one thread rename the file another is still + writing, and the loser dies on a missing path.""" path = os.path.join(config_dir, ".claude.json") try: with open(path, encoding="utf-8") as handle: @@ -151,8 +156,8 @@ def _seed_cli_identity(config_dir: str) -> None: return except (OSError, ValueError): pass - staged = f"{path}.{os.getpid()}" - with open(staged, "w", encoding="utf-8") as handle: + handle_fd, staged = tempfile.mkstemp(dir=config_dir, prefix=".claude.json.") + with os.fdopen(handle_fd, "w", encoding="utf-8") as handle: json.dump({"userID": _FIXED_CLI_USER_ID}, handle) os.replace(staged, path) From baca62df136bc7e7337b2bb3326e38aaf079dcac Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 13:53:28 -0700 Subject: [PATCH 124/187] test(logging): pick this test's own records out of the shared log batch The generic API logger batches whatever is queued when it flushes, so records from tests in other files in the same job land in the same request. Two tests assumed otherwise: one read actual_request[0], the other counted NDJSON lines, and both broke whenever another file logged first. Select by the messages each test sent instead, which keeps the format assertions and stops the order from deciding the outcome. --- .../test_generic_api_callback.py | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/tests/logging_callback_tests/test_generic_api_callback.py b/tests/logging_callback_tests/test_generic_api_callback.py index 29d8f9e5694..2c62741ed38 100644 --- a/tests/logging_callback_tests/test_generic_api_callback.py +++ b/tests/logging_callback_tests/test_generic_api_callback.py @@ -98,8 +98,15 @@ async def test_generic_api_callback(): assert isinstance(actual_request, list), "Request body should be a list" assert len(actual_request) > 0, "Request body list should not be empty" - # Validate the first payload item - payload_item: StandardLoggingPayload = StandardLoggingPayload(**actual_request[0]) + this_test_messages = [{"role": "user", "content": "Hello, world!"}] + mine = [ + item for item in actual_request if item.get("messages") == this_test_messages + ] + assert ( + len(mine) == 1 + ), f"Expected this test's single call in the batch, got {len(mine)} of {len(actual_request)}" + + payload_item: StandardLoggingPayload = StandardLoggingPayload(**mine[0]) print("##########\n") print(json.dumps(payload_item, indent=4)) print("##########\n") @@ -448,11 +455,15 @@ async def test_generic_api_callback_sumologic_uses_ndjson(): assert isinstance(ndjson_data, str), "Data should be a string for NDJSON" lines = ndjson_data.strip().split("\n") - assert len(lines) == 2, f"Expected 2 lines of NDJSON, got {len(lines)}" + records = [json.loads(line) for line in lines] - # Each line should be valid JSON - for line in lines: - json.loads(line) # Will raise if invalid JSON + this_test_messages = [ + [{"role": "user", "content": f"Test {i}"}] for i in range(2) + ] + mine = [record for record in records if record.get("messages") in this_test_messages] + assert ( + len(mine) == 2 + ), f"Expected this test's 2 calls as NDJSON lines, got {len(mine)} of {len(records)}" @pytest.mark.asyncio From b96a80400441aa4073c837e21d9c542a0aa3814e Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 20:55:32 +0000 Subject: [PATCH 125/187] fix(proxy): seed litellm_call_id into request data before parsing can fail The failure hook received data without the resolved id when body parsing or add_litellm_data_to_request raised, so proxy-only spend logging minted a fresh id that did not match the error log or the x-litellm-call-id header. The id is now part of the request data from the start and merged over the parsed body, which also removes the post-hoc in-place assignment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/image_endpoints/endpoints.py | 5 +- .../pass_through_endpoints.py | 7 ++- litellm/proxy/proxy_server.py | 15 +++--- litellm/proxy/rerank_endpoints/endpoints.py | 5 +- .../proxy/image_endpoints/test_endpoints.py | 51 ++++++++++++++++++- tests/test_litellm/proxy/test_proxy_server.py | 29 +++++++++++ 6 files changed, 92 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 30406bbcaae..5b90c0ff830 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -94,12 +94,12 @@ async def image_generation( version, ) - data = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data = {"litellm_call_id": litellm_call_id} try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() - data = orjson.loads(body) + data = orjson.loads(body) | data # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -110,7 +110,6 @@ async def image_generation( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id if isinstance(model, str): reject_url_valued_destination("model", model) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e93b1232836..c27d4f17016 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -199,15 +199,15 @@ async def chat_completion_pass_through_endpoint( version, ) - data = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data = {"litellm_call_id": litellm_call_id} try: body: Final = await request.body() body_str: Final = body.decode() try: - data = ast.literal_eval(body_str) + data = ast.literal_eval(body_str) | data except Exception: - data = json.loads(body_str) + data = json.loads(body_str) | data data["adapter_id"] = adapter_id @@ -228,7 +228,6 @@ async def chat_completion_pass_through_endpoint( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id # override with user settings, these are params passed via cli if user_temperature: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 45beb8cc93c..2fb239c0c92 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11471,12 +11471,12 @@ async def moderations( ``` """ global proxy_logging_obj - data: dict = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data: dict = {"litellm_call_id": litellm_call_id} try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() - data = orjson.loads(body) + data = orjson.loads(body) | data # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -11487,7 +11487,6 @@ async def moderations( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id data["model"] = ( general_settings.get("moderation_model", None) # server default @@ -11595,12 +11594,12 @@ async def audio_speech( https://platform.openai.com/docs/api-reference/audio/createSpeech """ global proxy_logging_obj - data: dict = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data: dict = {"litellm_call_id": litellm_call_id} try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() - data = orjson.loads(body) + data = orjson.loads(body) | data # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -11611,7 +11610,6 @@ async def audio_speech( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id if data.get("user", None) is None and user_api_key_dict.user_id is not None: data["user"] = user_api_key_dict.user_id @@ -11730,12 +11728,12 @@ async def audio_transcriptions( https://platform.openai.com/docs/api-reference/audio/createTranscription?lang=curl """ global proxy_logging_obj - data: dict = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data: dict = {"litellm_call_id": litellm_call_id} try: # Use orjson to parse JSON data, orjson speeds up requests significantly form_data: Final = await get_form_data(request) - data = {key: value for key, value in form_data.items() if key != "file"} + data = {key: value for key, value in form_data.items() if key != "file"} | data # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -11746,7 +11744,6 @@ async def audio_transcriptions( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id if data.get("user", None) is None and user_api_key_dict.user_id is not None: data["user"] = user_api_key_dict.user_id diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index 4f5eb411e44..4f2daed15ed 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -58,11 +58,11 @@ async def rerank( version, ) - data = {} litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) + data = {"litellm_call_id": litellm_call_id} try: body: Final = await request.body() - data = orjson.loads(body) + data = orjson.loads(body) | data # Include original request and headers in the data data = await add_litellm_data_to_request( @@ -73,7 +73,6 @@ async def rerank( version=version, proxy_config=proxy_config, ) - data["litellm_call_id"] = litellm_call_id ### CALL HOOKS ### - modify incoming data / reject request before calling the model data = await proxy_logging_obj.pre_call_hook(user_api_key_dict=user_api_key_dict, data=data, call_type="rerank") diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py index d03832bf6d0..31b87530c94 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py @@ -1,7 +1,7 @@ import asyncio import copy import logging -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from types import SimpleNamespace from typing import Any, Dict @@ -279,3 +279,52 @@ async def test_failure_log_carries_the_callers_litellm_call_id( record = next(r for r in caplog.records if "Exception occured" in r.getMessage()) assert record.litellm_call_id == call_id assert call_id in record.getMessage() + + +@pytest.mark.asyncio +async def test_failure_before_the_provider_call_bills_the_callers_litellm_call_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """LIT-7836: when the request is rejected while it is still being prepared, the + failure hook must see the same litellm_call_id the response header answers with, + otherwise the spend row is stored under a freshly minted id nobody can look up.""" + call_id = "images-early-7836" + hook_request_data: list[Mapping[str, object]] = [] + + async def rejecting_add_litellm_data_to_request(**_: object) -> object: + raise HTTPException(status_code=400, detail={"error": "tag not allowed"}) + + async def fake_post_call_failure_hook(*, request_data: Mapping[str, object], **_: object) -> None: + hook_request_data.append(request_data) + + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", rejecting_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", + SimpleNamespace(post_call_failure_hook=fake_post_call_failure_hook), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.version", "test-version") + + body = orjson.dumps({"model": "dall-e-3", "prompt": "a lighthouse at dusk", "litellm_call_id": "from-the-body"}) + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/images/generations", + "headers": [(b"x-litellm-call-id", call_id.encode())], + }, + receive, + ) + + with pytest.raises(ProxyException) as raised: + await endpoints.image_generation(request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()) + + assert raised.value.headers["x-litellm-call-id"] == call_id + assert [data["litellm_call_id"] for data in hook_request_data] == [call_id] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 14aea2bd020..a6e6a2100ca 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13046,6 +13046,35 @@ async def test_moderations_failure_log_carries_the_callers_litellm_call_id(caplo assert call_id in record.getMessage() +@pytest.mark.asyncio +async def test_moderations_unparseable_body_bills_the_callers_litellm_call_id(): + """LIT-7836: a body that fails to parse must still hand the failure hook the + litellm_call_id the response header answers with, so the spend row is findable.""" + from litellm.proxy._types import ProxyException + + call_id = "moderations-early-7836" + + request = MagicMock() + request.headers = {"x-litellm-call-id": call_id} + request.body = AsyncMock(return_value=b'{"input": ') + fake_logging = MagicMock() + fake_logging.post_call_failure_hook = AsyncMock() + + with ( + patch.object(proxy_server_module, "proxy_logging_obj", new=fake_logging), # test-quality-ok: module global, no injection point + pytest.raises(ProxyException) as raised, + ): + await proxy_server_module.moderations( + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", spend=0.0), + ) + + assert raised.value.headers["x-litellm-call-id"] == call_id + hook_request_data = fake_logging.post_call_failure_hook.await_args.kwargs["request_data"] + assert hook_request_data["litellm_call_id"] == call_id + + @pytest.mark.asyncio async def test_moderations_already_shaped_failure_answers_with_the_callers_litellm_call_id(): """LIT-7836: a ProxyException raised inside /v1/moderations is re-raised unwrapped but still From f10d95fb95371b1cebbec3ccd9fa3cce71fdc8ff Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 13:58:38 -0700 Subject: [PATCH 126/187] test(together_ai): drop the prefix-strip assertion, e2e covers it live tests/e2e/llm_translation/test_together_ai_e2e.py registers its model with the full registry key, so a slashed together_ai// goes through the prefix strip on every e2e run and an over-strip would fail against the real API. The unit assertion was a second copy of that. The roles check stays, since nothing in e2e exercises it. --- .../chat/test_together_ai_chat_transformation.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py index a7347edb2c7..1df8c96fb50 100644 --- a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py +++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py @@ -1137,21 +1137,6 @@ def _together_chat_transport() -> tuple[HTTPHandler, list[httpx.Request]]: return client, captured_requests -def test_only_the_provider_prefix_is_stripped_from_a_slashed_model_name(): - client, captured_requests = _together_chat_transport() - - litellm.completion( - model=f"together_ai/{TOOL_CALLING_MODEL}", - messages=[{"role": "user", "content": "Hello!"}], - api_key="fake-key", - client=client, - ) - - assert "/" in TOOL_CALLING_MODEL - assert str(captured_requests[0].url) == "https://api.together.ai/v1/chat/completions" - assert json.loads(captured_requests[0].content)["model"] == TOOL_CALLING_MODEL - - def test_custom_role_wrappers_never_reach_the_request(): client, captured_requests = _together_chat_transport() messages = [{"role": "user", "content": "Hello!"}] From c621435ef7de4178b8da74cbde5386400215bf2f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 21:02:30 +0000 Subject: [PATCH 127/187] refactor(ocr): move file preparation from the python bridge into litellm-core Delete litellm/ocr/input.py and the native _ocr_file_document, _ocr_upload_document and _ocr_mime_type helpers. File documents now project to a typed OcrDocumentInput and the core lifecycle reads local paths, encodes bytes and asks the host to read file-like objects through a ReadDocument operation before the provider request Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/core/src/ocr/client.rs | 13 +- litellm-rust/crates/core/src/ocr/document.rs | 126 ++++-- litellm-rust/crates/core/src/ocr/error.rs | 6 + litellm-rust/crates/core/src/ocr/lifecycle.rs | 50 ++- litellm-rust/crates/core/src/ocr/mod.rs | 7 +- litellm-rust/crates/core/src/ocr/types.rs | 70 +++- litellm-rust/crates/core/src/ocr/wire.rs | 43 +- litellm-rust/crates/core/tests/ocr.rs | 154 +++++++- litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../python-bridge/src/routes/ocr/document.rs | 370 ++++++++++-------- .../python-bridge/src/routes/ocr/errors.rs | 7 + .../python-bridge/src/routes/ocr/lifecycle.rs | 22 +- .../python-bridge/src/routes/ocr/mod.rs | 1 - .../python-bridge/src/routes/ocr/project.rs | 141 +++---- litellm/ocr/input.py | 112 ------ litellm/ocr/legacy.py | 7 +- litellm/ocr/main.py | 2 +- litellm/proxy/ocr_endpoints/endpoints.py | 17 +- litellm/rust_bridge/_native.pyi | 13 - tests/test_litellm/ocr/test_ocr_file_input.py | 25 +- tests/test_litellm_rust/ocr/test_requests.py | 139 +++---- 22 files changed, 810 insertions(+), 517 deletions(-) delete mode 100644 litellm/ocr/input.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e2a3af77594..7397742369b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1951,6 +1951,7 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "bytes", "criterion", "futures-util", "litellm-auth", diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 00bfeb2b7b2..9a30b2f8e04 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -54,9 +54,16 @@ impl OcrClient { match call.resume(result.take()).await? { OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().ok_or_else(|| { - Error::InvalidRequest("OCR request was already projected".into()) - })?), + Box::new( + request + .take() + .ok_or_else(|| { + Error::InvalidRequest( + "OCR request was already projected".into(), + ) + })? + .into(), + ), false, )))) } diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index a7afdaf8793..1b3d2dada44 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -1,3 +1,6 @@ +use std::io::Read; +use std::path::Path; + use base64::{Engine, engine::general_purpose::STANDARD}; use data_url::mime::Mime; use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError}; @@ -5,12 +8,52 @@ use reqwest::Url; use serde_json::Map; use super::error::{OcrError, OcrRequestError, OcrResponseError}; -use super::types::{OcrConnection, OcrDocument}; +use super::types::{OcrConnection, OcrDocument, OcrDocumentInput}; use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS}; use crate::media::Error as MediaError; use crate::media::{DownloadPolicy, MediaFetcher}; use crate::transport::Error as TransportError; +pub fn prepare_document(input: OcrDocumentInput) -> Result { + match input { + OcrDocumentInput::Document(document) => Ok(document), + OcrDocumentInput::Path { path, mime_type } => { + read_path_document(&path, mime_type.as_deref()) + } + OcrDocumentInput::Bytes { + bytes, + file_name, + mime_type, + } => Ok(encode_file_document( + &bytes, + file_name.as_deref(), + mime_type.as_deref(), + )?), + OcrDocumentInput::HostReader { .. } => Err(super::Error::InvalidRequest( + "OCR file reader was not read by the host".into(), + )), + } +} + +pub fn read_path_document( + path: &Path, + mime_type: Option<&str>, +) -> Result { + let mut bytes = Vec::new(); + std::fs::File::open(path) + .and_then(|file| { + file.take(OCR_INLINE_MAX_BYTES as u64 + 1) + .read_to_end(&mut bytes) + }) + .map_err(|source| super::Error::FileRead { + path: path.to_owned(), + kind: source.kind(), + message: source.to_string(), + })?; + let name = path.file_name().map(|name| name.to_string_lossy()); + Ok(encode_file_document(&bytes, name.as_deref(), mime_type)?) +} + pub fn encode_file_document( bytes: &[u8], file_name: Option<&str>, @@ -75,18 +118,6 @@ pub fn mime_type_for_name(name: &str) -> &'static str { } } -pub fn upload_mime_type<'a>(file_name: Option<&str>, content_type: Option<&'a str>) -> &'a str { - match content_type - .and_then(|value| value.split(';').next()) - .map(str::trim) - { - Some(value) if !value.is_empty() && value != "application/octet-stream" => value, - _ => file_name - .map(mime_type_for_name) - .unwrap_or("application/octet-stream"), - } -} - pub(crate) struct InlineDocument<'a>(DataUrl<'a>); impl<'a> InlineDocument<'a> { @@ -230,24 +261,65 @@ mod tests { } #[test] - fn upload_mime_mapping_matches_python() { + fn path_documents_are_read_and_named_by_core() { + let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("scan.png"); + std::fs::write(&path, b"abc").unwrap(); assert_eq!( - upload_mime_type(Some("report.pdf"), Some("application/octet-stream")), - "application/pdf" - ); - assert_eq!(upload_mime_type(Some("image.png"), None), "image/png"); - assert_eq!(upload_mime_type(None, None), "application/octet-stream"); - assert_eq!( - upload_mime_type(Some("doc.pdf"), Some("application/pdf; charset=utf-8")), - "application/pdf" + prepare_document(OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }) + .unwrap(), + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,YWJj".into(), + extra_fields: Map::new(), + } ); assert_eq!( - upload_mime_type( - Some("img.png"), - Some("image/png; charset=utf-8; boundary=something") - ), - "image/png" + prepare_document(OcrDocumentInput::Path { + path: path.clone(), + mime_type: Some("application/pdf".into()), + }) + .unwrap(), + document("data:application/pdf;base64,YWJj") ); + std::fs::write(&path, vec![b'a'; OCR_INLINE_MAX_BYTES + 1]).unwrap(); + assert_eq!( + prepare_document(OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }), + Err(OcrRequestError::InlineDocumentTooLarge.into()) + ); + std::fs::remove_dir_all(&dir).unwrap(); + + let missing = dir.join("missing.pdf"); + let Err(super::super::Error::FileRead { path, kind, .. }) = + prepare_document(OcrDocumentInput::Path { + path: missing.clone(), + mime_type: None, + }) + else { + panic!("missing paths must surface a file read error"); + }; + assert_eq!(path, missing); + assert_eq!(kind, std::io::ErrorKind::NotFound); + } + + #[test] + fn byte_documents_are_encoded_and_host_readers_must_be_read_first() { + assert_eq!( + prepare_document(OcrDocumentInput::Bytes { + bytes: b"abc".as_slice().into(), + file_name: Some("scan.pdf".into()), + mime_type: None, + }) + .unwrap(), + document("data:application/pdf;base64,YWJj") + ); + assert!(prepare_document(OcrDocumentInput::HostReader { mime_type: None }).is_err()); } #[test] diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs index 1c21edb6c91..0c92b511a38 100644 --- a/litellm-rust/crates/core/src/ocr/error.rs +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -50,6 +50,12 @@ pub enum Error { Connect(String), #[error("routing error: {0}")] Routing(String), + #[error("Failed to read OCR file {}: {message}", path.display())] + FileRead { + path: std::path::PathBuf, + kind: std::io::ErrorKind, + message: String, + }, /// The request is outside the surface this route covers in Rust. Hosts that /// keep a reference implementation treat this as "fall back", not "fail". #[error("unsupported by the rust path: {0}")] diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs index efa2b1f2873..994a9698459 100644 --- a/litellm-rust/crates/core/src/ocr/lifecycle.rs +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -9,6 +9,7 @@ use super::hooks::{ OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest, OcrPreCallRequest, }; +use super::types::{OcrDocumentInput, OcrFileContent}; use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient}; use crate::call_lifecycle::host::{ HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase, @@ -52,6 +53,7 @@ impl OcrAdmission { #[derive(Clone, Debug)] pub enum OcrHostOperation { ProjectRequest, + ReadDocument, Lifecycle(HostPhase), ConstructResponse(Arc), MapFailure(Error), @@ -83,7 +85,8 @@ impl OcrHostOperation { } pub enum OcrHostResult { - Request(Result<(Box, bool), Error>), + Request(Result<(Box>, bool), Error>), + Document(Result), Lifecycle(Result<(), HostFailure>), AzureAdToken(Result), PreCall(Result), @@ -313,7 +316,7 @@ struct PendingOperation { struct OcrExecution { client: Option, - request: Option, + request: Option>, operations_tx: mpsc::UnboundedSender, operations_rx: mpsc::UnboundedReceiver, pending_result: Option>, @@ -397,12 +400,14 @@ impl OcrExecution { }, ))); } - request.hooks = Arc::new(ProtocolHooks { + let hooks = Arc::new(ProtocolHooks { operations: self.operations_tx.clone(), intercepts_requests, terminal: self.terminal.clone(), }); + request.hooks = hooks.clone(); self.execution = Some(tokio::spawn(async move { + let request = prepare_request_document(request, &hooks).await?; perform_ocr_request(&client, request).await })); } @@ -423,6 +428,39 @@ impl OcrExecution { } } +async fn prepare_request_document( + request: LiteLLMOcrRequest, + hooks: &ProtocolHooks, +) -> Result { + let request = match &request.document { + OcrDocumentInput::HostReader { mime_type } => { + let mime_type = mime_type.clone(); + let content = match hooks.invoke(OcrHostOperation::ReadDocument).await? { + OcrHostResult::Document(result) => result?, + _ => { + return Err(Error::InvalidRequest( + "invalid OCR document read host result".into(), + )); + } + }; + request.with_document(OcrDocumentInput::Bytes { + bytes: content.bytes, + file_name: content.file_name, + mime_type, + }) + } + _ => request, + }; + if let OcrDocumentInput::Document(_) = &request.document { + return request.map_document(super::document::prepare_document); + } + tokio::task::spawn_blocking(move || request.map_document(super::document::prepare_document)) + .await + .map_err(|error| { + Error::InvalidRequest(format!("OCR document preparation task failed: {error}")) + })? +} + impl Drop for OcrExecution { fn drop(&mut self) { if let Some(execution) = &self.execution { @@ -567,6 +605,9 @@ impl OcrHost for NoopOcrHost { OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err( Error::InvalidRequest("OCR host has no request projection".into()), )), + OcrHostOperation::ReadDocument => OcrHostResult::Document(Err( + Error::InvalidRequest("OCR host has no document reader".into()), + )), OcrHostOperation::Lifecycle(_) | OcrHostOperation::ConstructResponse(_) | OcrHostOperation::MapFailure(_) @@ -602,6 +643,9 @@ impl OcrHost for OcrHookHost { OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err( Error::InvalidRequest("OCR hook host has no request projection".into()), )), + OcrHostOperation::ReadDocument => OcrHostResult::Document(Err( + Error::InvalidRequest("OCR hook host has no document reader".into()), + )), OcrHostOperation::Success { context, response, diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 3b51ff98356..f2e7aa4f46d 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -13,12 +13,15 @@ pub mod types; pub mod wire; pub use client::{OcrClient, ocr}; -pub use document::{encode_file_document, mime_type_for_name, upload_mime_type}; +pub use document::{encode_file_document, mime_type_for_name, read_path_document}; pub use lifecycle::{ NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult, }; -pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument}; +pub use types::{ + LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrDocumentInput, + OcrFileContent, +}; #[cfg(test)] #[path = "../../tests/azure_ai_ocr.rs"] diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 69e6982414b..bb212674b33 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,7 +1,10 @@ use std::collections::BTreeMap; +use std::convert::Infallible; +use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; +use bytes::Bytes; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -50,6 +53,35 @@ impl OcrDocument { } } +#[derive(Clone, Debug, PartialEq)] +pub enum OcrDocumentInput { + Document(OcrDocument), + Path { + path: PathBuf, + mime_type: Option, + }, + Bytes { + bytes: Bytes, + file_name: Option, + mime_type: Option, + }, + HostReader { + mime_type: Option, + }, +} + +impl From for OcrDocumentInput { + fn from(document: OcrDocument) -> Self { + Self::Document(document) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OcrFileContent { + pub bytes: Bytes, + pub file_name: Option, +} + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum OcrResponseFormat { @@ -89,9 +121,9 @@ impl Default for OcrConnection { } } -pub struct LiteLLMOcrRequest { +pub struct LiteLLMOcrRequest { pub model: String, - pub document: OcrDocument, + pub document: D, pub connection: OcrConnection, pub hooks: Arc, pub litellm_call_id: Option, @@ -101,10 +133,10 @@ pub struct LiteLLMOcrRequest { pub(crate) adapter: OcrAdapterKind, } -impl LiteLLMOcrRequest { +impl LiteLLMOcrRequest { pub fn new( model: String, - document: OcrDocument, + document: D, custom_llm_provider: Option<&str>, optional_params: Map, ) -> Result { @@ -151,6 +183,36 @@ impl LiteLLMOcrRequest { ..self } } + + pub fn map_document( + self, + map: impl FnOnce(D) -> Result, + ) -> Result, E> { + Ok(LiteLLMOcrRequest { + model: self.model, + document: map(self.document)?, + connection: self.connection, + hooks: self.hooks, + litellm_call_id: self.litellm_call_id, + optional_params: self.optional_params, + input_sources: self.input_sources, + azure_ad_token_provider: self.azure_ad_token_provider, + adapter: self.adapter, + }) + } + + pub fn with_document(self, document: T) -> LiteLLMOcrRequest { + let Ok(request) = self.map_document(|_| Ok::(document)); + request + } +} + +impl From for LiteLLMOcrRequest { + fn from(request: LiteLLMOcrRequest) -> Self { + let Ok(request) = request + .map_document(|document| Ok::<_, Infallible>(OcrDocumentInput::Document(document))); + request + } } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 93816effcb1..f0cad2b4e93 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -68,9 +68,9 @@ pub struct DecodedOcrResponse { #[derive(Deserialize)] #[serde(deny_unknown_fields)] -pub struct OcrWireRequest { +pub struct OcrWireRequest { pub model: String, - pub document: Value, + pub document: D, pub api_key: Option, pub api_base: Option, pub custom_llm_provider: Option, @@ -141,10 +141,34 @@ pub fn consumed_optional_params( } pub fn decode_request(wire: OcrWireRequest) -> Result { + let OcrWireRequest { + model, + document, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + input_sources, + timeout_seconds, + } = wire; + decode_request_input(OcrWireRequest { + model, + document: decode_document(document)?, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + input_sources, + timeout_seconds, + }) +} + +pub fn decode_request_input(wire: OcrWireRequest) -> Result, Error> { let api_key_source = source_for(&wire.input_sources, "api_key"); let api_base_source = source_for(&wire.input_sources, "api_base"); let extra_headers_source = source_for(&wire.input_sources, "extra_headers"); - let document = decode_document(wire.document)?; let headers = wire .extra_headers .unwrap_or_default() @@ -183,7 +207,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result .unwrap_or(defaults.max_response_bytes); let request = LiteLLMOcrRequest::new( wire.model, - document, + wire.document, wire.custom_llm_provider.as_deref(), wire.optional_params .into_iter() @@ -209,14 +233,14 @@ pub fn decode_request(wire: OcrWireRequest) -> Result }) } -fn decode_document(value: Value) -> Result { +pub fn decode_document(value: Value) -> Result { let kind = value.get("type").and_then(Value::as_str); let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none() || matches!(kind, Some("image_url")) && value.get("image_url").is_none(); if missing_url { - return Err(OcrRequestError::MissingDocumentUrl); + return Err(OcrRequestError::MissingDocumentUrl.into()); } - decode_request_value(value, "document") + Ok(decode_request_value(value, "document")?) } fn source_for(sources: &BTreeMap, name: &str) -> InputSource { @@ -334,10 +358,7 @@ mod tests { serde_json::json!({"type": "document_url"}), serde_json::json!({"type": "image_url"}), ] { - assert_eq!( - decode_document(document), - Err(OcrRequestError::MissingDocumentUrl) - ); + assert_eq!(decode_document(document), Err(Error::MissingDocumentUrl)); } } } diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 373972cf68b..a24d960422d 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -348,13 +348,14 @@ async fn fallible_host_phases_do_not_replay_or_reach_transport() { } OcrHostOperation::ProjectRequest => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap()), + Box::new(request.take().unwrap().into()), false, )))) } OcrHostOperation::AcquireAzureAdToken => { panic!("test request has no token provider") } + OcrHostOperation::ReadDocument => panic!("test request has no file reader"), OcrHostOperation::PreCall(request) => { phases.push("pre"); result = Some(OcrHostResult::PreCall(if failure_phase == "pre" { @@ -405,7 +406,7 @@ async fn invalid_provider_response_runs_post_call_before_normalization_failure() match call.resume(result.take()).await { Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap()), + Box::new(request.take().unwrap().into()), false, )))); } @@ -467,9 +468,10 @@ async fn direct_native_host_drives_the_same_state_machine() { _ => panic!("unexpected OCR operation"), }); result = Some(match operation { - OcrHostOperation::ProjectRequest => { - OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) - } + OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap().into()), + false, + ))), operation => host.invoke(operation).await, }); } @@ -501,6 +503,137 @@ async fn direct_native_host_drives_the_same_state_machine() { )); } +async fn drive_native_file_call( + request: super::LiteLLMOcrRequest, + content: Result, +) -> (Result, usize) { + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let mut content = Some(content); + let mut result = None; + let mut reads = 0; + let outcome = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))); + } + Ok(OcrCallStep::Host(OcrHostOperation::ReadDocument)) => { + reads += 1; + result = Some(OcrHostResult::Document(content.take().unwrap())); + } + Ok(OcrCallStep::Host(operation)) => result = Some(NoopOcrHost.invoke(operation).await), + Ok(OcrCallStep::Complete(response)) => break Ok(response), + Err(error) => break Err(error), + } + }; + (outcome, reads) +} + +#[tokio::test] +async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"file"}] + }))]) + .await; + let request = wire_request("mistral/model", &base, json!({})).with_document( + super::OcrDocumentInput::HostReader { + mime_type: Some("application/pdf".into()), + }, + ); + let (response, reads) = drive_native_file_call( + request, + Ok(super::OcrFileContent { + bytes: b"abc".as_slice().into(), + file_name: Some("scan.png".into()), + }), + ) + .await; + server.await.unwrap(); + assert_eq!(response.unwrap().pages[0]["markdown"], "file"); + assert_eq!(reads, 1); + assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj")); +} + +#[tokio::test] +async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() { + let (base, seen, _server) = mock_server(vec![]).await; + let request = wire_request("mistral/model", &base, json!({})); + let failure = crate::ocr::Error::InvalidRequest("reader exploded".into()); + let (response, reads) = drive_native_file_call( + request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }), + Err(failure.clone()), + ) + .await; + assert_eq!(response.unwrap_err(), failure); + assert_eq!(reads, 1); + + let request = wire_request("mistral/model", &base, json!({})); + let (response, _) = drive_native_file_call( + request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }), + Ok(super::OcrFileContent { + bytes: Default::default(), + file_name: None, + }), + ) + .await; + assert!(matches!( + response.unwrap_err(), + crate::ocr::Error::InvalidRequest(_) + )); + assert!(seen.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn path_documents_are_read_by_core_without_a_host_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"path"}] + }))]) + .await; + let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("scan.png"); + std::fs::write(&path, b"abc").unwrap(); + let request = wire_request("mistral/model", &base, json!({})).with_document( + super::OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }, + ); + let (response, reads) = drive_native_file_call( + request, + Err(crate::ocr::Error::InvalidRequest("unused".into())), + ) + .await; + server.await.unwrap(); + std::fs::remove_dir_all(&dir).unwrap(); + assert_eq!(response.unwrap().pages[0]["markdown"], "path"); + assert_eq!(reads, 0); + assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); + + let (base, seen, _server) = mock_server(vec![]).await; + let request = wire_request("mistral/model", &base, json!({})); + let (response, _) = drive_native_file_call( + request.with_document(super::OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }), + Err(crate::ocr::Error::InvalidRequest("unused".into())), + ) + .await; + assert!(matches!( + response.unwrap_err(), + crate::ocr::Error::FileRead { path: failed, kind: std::io::ErrorKind::NotFound, .. } if failed == path + )); + assert!(seen.lock().unwrap().is_empty()); +} + #[tokio::test] async fn public_finalization_failure_never_dispatches_success_or_replays_provider() { use crate::call_lifecycle::host::{HostFailure, HostPhase}; @@ -543,9 +676,10 @@ async fn public_finalization_failure_never_dispatches_success_or_replays_provide | OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => { panic!("finalization failure used provider/success dispatch") } - OcrHostOperation::ProjectRequest => { - OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) - } + OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap().into()), + false, + ))), operation => host.invoke(operation).await, }); } @@ -582,7 +716,7 @@ async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break, OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap()), + Box::new(request.take().unwrap().into()), false, )))) } @@ -799,7 +933,7 @@ async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_ _ = entered.notified() => break, step = call.resume(result.take()) => { result = Some(match step.unwrap() { - OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))), + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap().into()), false))), OcrCallStep::Host(operation) => NoopOcrHost.invoke(operation).await, OcrCallStep::Complete(_) => panic!("pending provider completed"), }); diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1562d4c1021..6dde7c71af6 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -16,6 +16,7 @@ extension-module = ["pyo3/extension-module"] panic-test = [] [dependencies] +bytes.workspace = true futures-util.workspace = true litellm-core.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs index d43c2f88775..33c0561184d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs @@ -1,97 +1,56 @@ -use std::io::Read; use std::path::PathBuf; -use pyo3::exceptions::{PyFileNotFoundError, PyTypeError, PyValueError}; +use bytes::Bytes; +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use pyo3::pybacked::PyBackedBytes; -#[cfg(test)] -use pyo3::types::PyDict; use pyo3::types::{PyBytes, PyString}; -use litellm_core::constants::OCR_INLINE_MAX_BYTES; -use litellm_core::ocr::{OcrDocument, encode_file_document, mime_type_for_name, upload_mime_type}; -use litellm_python_interop::to_py_preserving_errors; +use litellm_core::ocr::{OcrDocumentInput, OcrFileContent}; -enum FileBytes { - Python(PyBackedBytes), - Native(Vec), +#[derive(Debug)] +pub(super) struct PythonFileReader { + reader: Py, + name: Option, } -impl AsRef<[u8]> for FileBytes { - fn as_ref(&self) -> &[u8] { - match self { - Self::Python(bytes) => bytes, - Self::Native(bytes) => bytes, - } +impl PythonFileReader { + pub(super) fn read(&self, py: Python<'_>) -> PyResult { + let value = self.reader.bind(py).call0()?; + let bytes = if value.is_instance_of::() { + Bytes::from(value.extract::()?) + } else if value.is_instance_of::() { + extract_bytes(&value)? + } else { + return Err(PyTypeError::new_err(format!( + "OCR file read must return bytes or str, got {}", + value.get_type(), + ))); + }; + Ok(OcrFileContent { + bytes, + file_name: self.name.clone(), + }) + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reader) } } -fn read_file_input( - py: Python<'_>, - file: &Bound<'_, PyAny>, -) -> PyResult<(FileBytes, Option)> { - if file.is_instance_of::() { - return Err(PyValueError::new_err( - "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.", - )); +fn extract_bytes(value: &Bound<'_, PyAny>) -> PyResult { + if value.is_exact_instance_of::() { + return Ok(Bytes::from_owner(value.extract::()?)); } - if file.is_instance(&py.import("os")?.getattr("PathLike")?)? { - let path: PathBuf = file.extract()?; - let name = path - .file_name() - .map(|value| value.to_string_lossy().into_owned()); - let bytes = py - .detach(|| { - let mut bytes = Vec::new(); - std::fs::File::open(&path)? - .take(OCR_INLINE_MAX_BYTES as u64 + 1) - .read_to_end(&mut bytes)?; - Ok::<_, std::io::Error>(bytes) - }) - .map_err(|error| { - if error.kind() == std::io::ErrorKind::NotFound { - PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) - } else { - error.into() - } - })?; - return Ok((FileBytes::Native(bytes), name)); - } - if file.is_instance_of::() { - return Ok((FileBytes::Python(file.extract()?), None)); - } - let reader = file - .getattr_opt("read")? - .filter(|value| value.is_callable()); - let Some(reader) = reader else { - return Err(PyValueError::new_err(format!( - "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.", - file.get_type(), - ))); - }; - let name = file - .getattr_opt("name")? - .filter(|value| !value.is_none()) - .map(|value| value.extract::()) - .transpose()?; - let value = reader.call0()?; - let bytes = if value.is_instance_of::() { - FileBytes::Native(value.extract::()?.into_bytes()) - } else if value.is_instance_of::() { - FileBytes::Python(value.extract()?) - } else { - return Err(PyTypeError::new_err(format!( - "OCR file read must return bytes or str, got {}", - value.get_type(), - ))); - }; - Ok((bytes, name)) + Ok(Bytes::copy_from_slice( + value.extract::()?.as_ref(), + )) } pub(super) struct FileDocumentInput { - bytes: FileBytes, - name: Option, - mime_type: Option, + pub input: OcrDocumentInput, + pub reader: Option, } impl FromPyObject<'_, '_> for FileDocumentInput { @@ -104,79 +63,79 @@ impl FromPyObject<'_, '_> for FileDocumentInput { Err(error) if error.is_instance_of::(py) => None, Err(error) => return Err(error), }; + let missing = || { + PyValueError::new_err( + "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes", + ) + }; let file = document.get_item("file").map_err(|error| { if error.is_instance_of::(py) { - PyValueError::new_err("document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes") + missing() } else { error } })?; if file.is_none() { + return Err(missing()); + } + if file.is_instance_of::() { return Err(PyValueError::new_err( - "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes", + "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.", )); } - let (bytes, name) = read_file_input(py, &file)?; + if file.is_instance(&py.import("os")?.getattr("PathLike")?)? { + return Ok(Self { + input: OcrDocumentInput::Path { + path: file.extract::()?, + mime_type, + }, + reader: None, + }); + } + if file.is_instance_of::() { + return Ok(Self { + input: OcrDocumentInput::Bytes { + bytes: extract_bytes(&file)?, + file_name: None, + mime_type, + }, + reader: None, + }); + } + let reader = file + .getattr_opt("read")? + .filter(|value| value.is_callable()); + let Some(reader) = reader else { + return Err(PyValueError::new_err(format!( + "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.", + file.get_type(), + ))); + }; + let name = file + .getattr_opt("name")? + .filter(|value| !value.is_none()) + .map(|value| value.extract::()) + .transpose()?; Ok(Self { - bytes, - name, - mime_type, + input: OcrDocumentInput::HostReader { mime_type }, + reader: Some(PythonFileReader { + reader: reader.unbind(), + name, + }), }) } } -pub(super) fn file_document(py: Python<'_>, document: FileDocumentInput) -> PyResult { - py.detach(|| { - encode_file_document( - document.bytes.as_ref(), - document.name.as_deref(), - document.mime_type.as_deref(), - ) - }) - .map_err(|error| PyValueError::new_err(error.to_string())) -} - -#[pyfunction] -fn _ocr_file_document(py: Python<'_>, document: Bound<'_, PyAny>) -> PyResult> { - to_py_preserving_errors(py, &file_document(py, document.extract()?)?) -} - -#[pyfunction] -fn _ocr_mime_type(file_name: &str) -> String { - mime_type_for_name(file_name).into() -} - -#[pyfunction] -#[pyo3(signature = (file_content, file_name=None, content_type=None))] -fn _ocr_upload_document( - py: Python<'_>, - file_content: &Bound<'_, PyBytes>, - file_name: Option<&str>, - content_type: Option<&str>, -) -> PyResult> { - let bytes: PyBackedBytes = file_content.extract()?; - let document = py - .detach(|| { - encode_file_document( - &bytes, - None, - Some(upload_mime_type(file_name, content_type)), - ) - }) - .map_err(|error| PyValueError::new_err(error.to_string()))?; - to_py_preserving_errors(py, &document) -} - -pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add("_OCR_MAX_FILE_BYTES", OCR_INLINE_MAX_BYTES)?; - module.add_function(wrap_pyfunction!(_ocr_upload_document, module)?)?; - module.add_function(wrap_pyfunction!(_ocr_file_document, module)?)?; - module.add_function(wrap_pyfunction!(_ocr_mime_type, module)?) -} - #[cfg(test)] mod tests { use super::*; + use pyo3::types::PyDict; + + fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + locals + } #[test] fn extraction_validates_required_file_and_optional_mime_type() { @@ -196,69 +155,148 @@ mod tests { let error = document.extract::().err().unwrap(); assert!(error.is_instance_of::(py)); } - let document = py.eval(c"{'file': b'abc'}", None, None).unwrap(); + let error = py + .eval(c"{'file': 'scan.pdf'}", None, None) + .unwrap() + .extract::() + .err() + .unwrap(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("bare str")); + let document = py + .eval(c"{'file': b'abc', 'mime_type': 'image/png'}", None, None) + .unwrap(); let input: FileDocumentInput = document.extract().unwrap(); - assert_eq!(input.bytes.as_ref(), b"abc"); - assert_eq!(input.name, None); - assert_eq!(input.mime_type, None); + assert!(input.reader.is_none()); + assert_eq!( + input.input, + OcrDocumentInput::Bytes { + bytes: b"abc".as_slice().into(), + file_name: None, + mime_type: Some("image/png".into()), + } + ); }); } #[test] - fn extraction_validates_mime_type_before_consuming_file() { + fn paths_and_readers_are_projected_without_io() { Python::initialize(); Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - c"class Reader: + let locals = eval( + py, + c"from pathlib import Path +class Reader: + name = 'scan.png' def __init__(self): self.reads = 0 def read(self): self.reads += 1 return b'abc' reader = Reader() -document = {'file': reader, 'mime_type': 7}", - Some(&locals), - Some(&locals), - ) - .unwrap(); +document = {'file': reader, 'mime_type': 7} +reader_document = {'file': reader} +path_document = {'file': Path('/nonexistent/ocr-projection-test.pdf'), 'mime_type': 'image/png'}", + ); let document = locals.get_item("document").unwrap().unwrap(); let error = document.extract::().err().unwrap(); assert!(error.is_instance_of::(py)); - let reads: usize = locals - .get_item("reader") - .unwrap() - .unwrap() - .getattr("reads") - .unwrap() - .extract() - .unwrap(); - assert_eq!(reads, 0); + + let document = locals.get_item("reader_document").unwrap().unwrap(); + let input: FileDocumentInput = document.extract().unwrap(); + assert_eq!( + input.input, + OcrDocumentInput::HostReader { mime_type: None } + ); + let reads = || { + locals + .get_item("reader") + .unwrap() + .unwrap() + .getattr("reads") + .unwrap() + .extract::() + .unwrap() + }; + assert_eq!(reads(), 0); + let content = input.reader.unwrap().read(py).unwrap(); + assert_eq!(reads(), 1); + assert_eq!( + content, + OcrFileContent { + bytes: b"abc".as_slice().into(), + file_name: Some("scan.png".into()), + } + ); + + let document = locals.get_item("path_document").unwrap().unwrap(); + let input: FileDocumentInput = document.extract().unwrap(); + assert!(input.reader.is_none()); + assert_eq!( + input.input, + OcrDocumentInput::Path { + path: PathBuf::from("/nonexistent/ocr-projection-test.pdf"), + mime_type: Some("image/png".into()), + } + ); }); } #[test] - fn extraction_preserves_reader_key_error_identity() { + fn reader_results_are_normalized_and_exceptions_keep_their_identity() { Python::initialize(); Python::attach(|py| { - let locals = PyDict::new(py); - py.run( + let locals = eval( + py, c"failure = KeyError('reader failed') -class Reader: +class Raising: def read(self): raise failure -document = {'file': Reader()}", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let document = locals.get_item("document").unwrap().unwrap(); - let error = document.extract::().err().unwrap(); +class Text: + def read(self): + return 'héllo' +class Wrong: + def read(self): + return 7 +raising = {'file': Raising()} +text = {'file': Text()} +wrong = {'file': Wrong()}", + ); + let reader = |name: &str| { + locals + .get_item(name) + .unwrap() + .unwrap() + .extract::() + .unwrap() + .reader + .unwrap() + }; + let error = reader("raising").read(py).unwrap_err(); assert!( error .value(py) .is(locals.get_item("failure").unwrap().unwrap()) ); + assert_eq!( + reader("text").read(py).unwrap().bytes.as_ref(), + "héllo".as_bytes() + ); + let error = reader("wrong").read(py).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("bytes or str")); }); } + + #[test] + fn exact_python_bytes_transfer_without_copying_and_outlive_the_input() { + Python::initialize(); + let (bytes, pointer) = Python::attach(|py| { + let value = PyBytes::new(py, b"document bytes"); + let pointer = value.as_bytes().as_ptr() as usize; + (extract_bytes(value.as_any()).unwrap(), pointer) + }); + assert_eq!(bytes.as_ptr() as usize, pointer); + assert_eq!(bytes.as_ref(), b"document bytes"); + } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs index e4ce813d297..7dbc35289ff 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -1,4 +1,5 @@ use litellm_core::ocr::Error; +use pyo3::exceptions::{PyFileNotFoundError, PyOSError}; use pyo3::prelude::*; use crate::errors::{RustUpstreamError, core_error_to_pyerr}; @@ -7,6 +8,12 @@ pub(super) fn to_pyerr(error: Error) -> PyErr { let status = error.http_status_code(); let mapped = match error { Error::Http { status, body } => RustUpstreamError::new_err((status, body)), + Error::FileRead { + path, + kind: std::io::ErrorKind::NotFound, + .. + } => PyFileNotFoundError::new_err(format!("File not found: {}", path.display())), + Error::FileRead { message, .. } => PyOSError::new_err(message), other => core_error_to_pyerr(other.into()), }; attach_status(mapped, status) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs index 32794936899..e710b0d82f9 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -66,13 +66,27 @@ impl PythonOcrHost { retained_fields.set_item(name, value)?; } } - retained_fields.set_item("document", &self.projected()?.fields.document)?; let projected = self.projected_mut()?; + let document = match &projected.fields.document { + Some(document) => document.clone_ref(py), + None => to_py(py, &request.document)?, + }; + retained_fields.set_item("document", &document)?; + projected.fields.document = Some(document); projected.retained_fields = Some(retained_fields.unbind()); projected.pre_call = Some((&request).into()); Ok(request) } + fn read_document(&self, py: Python<'_>) -> PyResult { + self.projected()? + .fields + .reader + .as_ref() + .ok_or_else(missing_state)? + .read(py) + } + fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult { let provider = self .projected()? @@ -193,7 +207,7 @@ impl PythonRoute for PythonOcrHost { let OcrHostData::Unprojected { request } = &self.data else { return Err(missing_state()); }; - let projected = project_request(py, request.bind(py), self.state.kwargs.bind(py))?; + let projected = project_request(request.bind(py), self.state.kwargs.bind(py))?; let has_token_provider = projected.fields.azure_ad_token_provider.is_some(); let request = projected.request; self.data = OcrHostData::Projected(Box::new(ProjectedOcrHost { @@ -205,6 +219,7 @@ impl PythonRoute for PythonOcrHost { })); OcrHostResult::Request(Ok((Box::new(request), has_token_provider))) } + OcrHostOperation::ReadDocument => OcrHostResult::Document(Ok(self.read_document(py)?)), OcrHostOperation::AcquireAzureAdToken => { OcrHostResult::AzureAdToken(Ok(self.acquire_azure_ad_token(py)?)) } @@ -258,6 +273,9 @@ impl PythonRoute for PythonOcrHost { OcrHostData::Projected(projected) => { visit.call(&projected.fields.boundary_request)?; visit.call(&projected.fields.document)?; + if let Some(reader) = &projected.fields.reader { + reader.traverse(visit)?; + } visit.call(&projected.fields.api_key)?; if let Some(provider) = &projected.fields.azure_ad_token_provider { provider.traverse(visit)?; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index f17bf249b7f..5eae8ccf33f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -9,6 +9,5 @@ use pyo3::prelude::*; pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { value::register(module)?; - document::register(module)?; lifecycle::register(module) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 8d8d5f8c518..ad223645c62 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,14 +1,15 @@ use std::sync::Arc; -use litellm_core::ocr::wire::{OcrWireRequest, consumed_optional_params, decode_request}; -use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall}; -use litellm_python_interop::{ - from_py_preserving_errors as from_py, to_py_preserving_errors as to_py, +use litellm_core::ocr::wire::{ + OcrWireRequest, consumed_optional_params, decode_document, decode_request_input, }; +use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall, OcrDocumentInput}; +use litellm_python_interop::from_py_preserving_errors as from_py; use pyo3::prelude::*; use pyo3::types::PyDict; use serde_json::{Map, Value}; +use super::document::{FileDocumentInput, PythonFileReader}; use super::errors::to_pyerr as ocr_error_to_pyerr; use super::lifecycle::BridgeOcrHooks; use crate::auth::{AZURE_AD_TOKEN_PROVIDER, PythonTokenProvider}; @@ -17,7 +18,8 @@ use crate::marshal::{project_optional_fields, python_timeout_seconds, request_in pub(super) struct ProjectedOcrFields { pub boundary_request: Py, - pub document: Py, + pub document: Option>, + pub reader: Option, pub api_key: Py, pub azure_ad_token_provider: Option, pub provider: &'static str, @@ -25,7 +27,7 @@ pub(super) struct ProjectedOcrFields { } pub(super) struct ProjectedOcrCall { - pub request: LiteLLMOcrRequest, + pub request: LiteLLMOcrRequest, pub fields: ProjectedOcrFields, } @@ -80,12 +82,12 @@ impl<'py> OcrArguments<'_, 'py> { } enum ProjectedDocument { - File { wire: Value, retained: Py }, + File(FileDocumentInput), Other { wire: Value, retained: Py }, } impl ProjectedDocument { - fn project(py: Python<'_>, document: &Bound<'_, PyAny>) -> PyResult { + fn project(document: &Bound<'_, PyAny>) -> PyResult { let kind: String = document.get_item("type")?.extract()?; if kind != "file" { return Ok(Self::Other { @@ -93,25 +95,28 @@ impl ProjectedDocument { retained: document.clone().unbind(), }); } - let input = document.extract()?; - let encoded = super::document::file_document(py, input)?; - let wire = serde_json::to_value(encoded) - .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?; - Ok(Self::File { - retained: to_py(py, &wire)?, - wire, - }) + Ok(Self::File(document.extract()?)) } - fn into_parts(self) -> (Value, Py) { + fn into_parts( + self, + ) -> PyResult<( + OcrDocumentInput, + Option>, + Option, + )> { match self { - Self::File { wire, retained } | Self::Other { wire, retained } => (wire, retained), + Self::File(FileDocumentInput { input, reader }) => Ok((input, None, reader)), + Self::Other { wire, retained } => Ok(( + decode_document(wire).map_err(ocr_error_to_pyerr)?.into(), + Some(retained), + None, + )), } } } pub(super) fn project_request( - py: Python<'_>, request: &Bound<'_, PyAny>, kwargs: &Bound<'_, PyDict>, ) -> PyResult { @@ -119,8 +124,7 @@ pub(super) fn project_request( let arguments = OcrArguments { request, kwargs }; let model = arguments.model()?; let custom_llm_provider = arguments.custom_llm_provider()?; - let (wire_document, retained_document) = - ProjectedDocument::project(py, &arguments.document()?)?.into_parts(); + let document = ProjectedDocument::project(&arguments.document()?)?; let api_key = arguments.api_key()?; let specs = consumed_optional_params(&model, custom_llm_provider.as_deref()) .map_err(ocr_error_to_pyerr)?; @@ -136,9 +140,10 @@ pub(super) fn project_request( let azure_ad_token_provider = kwargs .get_item("azure_ad_token_provider")? .and_then(|provider| PythonTokenProvider::select(provider, AZURE_AD_TOKEN_PROVIDER)); + let (document, retained_document, reader) = document.into_parts()?; let wire = OcrWireRequest { model, - document: wire_document, + document, api_key: api_key.extract()?, api_base: arguments.api_base()?, custom_llm_provider, @@ -147,13 +152,14 @@ pub(super) fn project_request( input_sources, timeout_seconds: arguments.timeout_seconds()?, }; - let request = decode_request(wire).map_err(ocr_error_to_pyerr)?; + let request = decode_request_input(wire).map_err(ocr_error_to_pyerr)?; let provider = request.provider_name(); Ok(ProjectedOcrCall { request: request.with_host_hooks(Arc::new(BridgeOcrHooks), None), fields: ProjectedOcrFields { boundary_request, document: retained_document, + reader, api_key: api_key.unbind(), azure_ad_token_provider, provider, @@ -197,10 +203,21 @@ mod tests { } fn project_document( - py: Python<'_>, document: &Bound<'_, PyAny>, - ) -> PyResult<(Value, Py)> { - ProjectedDocument::project(py, document).map(ProjectedDocument::into_parts) + ) -> PyResult<( + OcrDocumentInput, + Option>, + Option, + )> { + ProjectedDocument::project(document)?.into_parts() + } + + fn url_document(url: &str) -> OcrDocumentInput { + litellm_core::ocr::OcrDocument::DocumentUrl { + document_url: url.into(), + extra_fields: Map::new(), + } + .into() } fn stub_timeout_conversion(py: Python<'_>) { @@ -374,7 +391,7 @@ kwargs = {} } #[test] - fn document_reader_mutations_are_visible_to_later_field_reads() { + fn document_readers_are_not_consumed_during_projection() { Python::initialize(); Python::attach(|py| { stub_timeout_conversion(py); @@ -406,7 +423,12 @@ kwargs = {} .unwrap(); let arguments = arguments(&request, &kwargs); let document = arguments.document().unwrap(); - project_document(py, &document).unwrap(); + let (input, retained, reader) = project_document(&document).unwrap(); + assert_eq!(input, OcrDocumentInput::HostReader { mime_type: None }); + assert!(retained.is_none()); + assert_eq!(arguments.api_base().unwrap().as_deref(), Some("original")); + assert_eq!(arguments.timeout_seconds().unwrap(), Some(1.0)); + reader.unwrap().read(py).unwrap(); assert_eq!(arguments.api_base().unwrap().as_deref(), Some("mutated")); assert_eq!(arguments.timeout_seconds().unwrap(), Some(9.0)); }); @@ -444,7 +466,7 @@ kwargs = {'api_key': key} } #[test] - fn file_documents_are_encoded_and_other_documents_keep_the_python_object() { + fn file_documents_become_typed_inputs_and_other_documents_keep_the_python_object() { Python::initialize(); Python::attach(|py| { let file = py @@ -454,13 +476,17 @@ kwargs = {'api_key': key} None, ) .unwrap(); + let (input, retained, reader) = project_document(&file).unwrap(); assert_eq!( - project_document(py, &file).unwrap().0, - serde_json::json!({ - "type": "document_url", - "document_url": "data:application/pdf;base64,JVBERi0xLjQ=", - }) + input, + OcrDocumentInput::Bytes { + bytes: b"%PDF-1.4".as_slice().into(), + file_name: None, + mime_type: Some("application/pdf".into()), + } ); + assert!(retained.is_none()); + assert!(reader.is_none()); let original = py .eval( @@ -469,44 +495,21 @@ kwargs = {'api_key': key} None, ) .unwrap(); - let (wire, retained) = project_document(py, &original).unwrap(); - assert_eq!( - wire, - serde_json::json!({ - "type": "document_url", - "document_url": "https://example.com/a.pdf", - }) - ); - assert!(retained.bind(py).is(&original)); + let (input, retained, _) = project_document(&original).unwrap(); + assert_eq!(input, url_document("https://example.com/a.pdf")); + assert!(retained.unwrap().bind(py).is(&original)); }); } #[test] - fn unknown_document_types_reach_existing_downstream_validation() { + fn unknown_document_types_reach_existing_core_validation() { Python::initialize(); Python::attach(|py| { let document = py .eval(c"{'type': 'mystery', 'mystery': 'x'}", None, None) .unwrap(); - let wire_document = project_document(py, &document).unwrap().0; - assert_eq!( - wire_document, - serde_json::json!({"type": "mystery", "mystery": "x"}) - ); - let error = match decode_request(OcrWireRequest { - model: "mistral/mistral-ocr-latest".into(), - document: wire_document, - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Map::new(), - input_sources: Default::default(), - timeout_seconds: None, - }) { - Ok(_) => panic!("unknown discriminators belong to core validation"), - Err(error) => error, - }; + let error = project_document(&document).unwrap_err(); + assert!(error.is_instance_of::(py)); assert!(error.to_string().contains("document")); }); } @@ -517,14 +520,14 @@ kwargs = {'api_key': key} Python::attach(|py| { let missing = py.eval(c"{}", None, None).unwrap(); assert!( - project_document(py, &missing) + project_document(&missing) .unwrap_err() .is_instance_of::(py) ); let non_string = py.eval(c"{'type': 1}", None, None).unwrap(); assert!( - project_document(py, &non_string) + project_document(&non_string) .unwrap_err() .is_instance_of::(py) ); @@ -540,7 +543,7 @@ document = Document() ", ); let error = - project_document(py, &locals.get_item("document").unwrap().unwrap()).unwrap_err(); + project_document(&locals.get_item("document").unwrap().unwrap()).unwrap_err(); assert!( error .value(py) @@ -569,9 +572,9 @@ document = Document() ", ); let document = locals.get_item("document").unwrap().unwrap(); - let (wire, retained) = project_document(py, &document).unwrap(); - assert_eq!(wire["type"], "document_url"); - assert!(!retained.bind(py).is(&document)); + let (input, retained, _) = project_document(&document).unwrap(); + assert!(matches!(input, OcrDocumentInput::Bytes { .. })); + assert!(retained.is_none()); let reads: Vec = document.getattr("reads").unwrap().extract().unwrap(); assert_eq!(reads, ["type", "mime_type", "file"]); }); diff --git a/litellm/ocr/input.py b/litellm/ocr/input.py deleted file mode 100644 index bcb448371c4..00000000000 --- a/litellm/ocr/input.py +++ /dev/null @@ -1,112 +0,0 @@ -from collections.abc import Mapping -from os import PathLike -from typing import Final, Literal, Protocol, cast # noqa: TID251 # native callables are validated when loaded - -from typing_extensions import NotRequired, ReadOnly, TypedDict - -from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.configuration import rust_ocr_enabled - - -class FileReader(Protocol): - def read(self) -> bytes | str: ... - - -class FileDocument(TypedDict): - type: ReadOnly[Literal["file"]] - file: ReadOnly[bytes | PathLike[str] | FileReader] - mime_type: ReadOnly[NotRequired[str]] - - -class NativeFileDocument(Protocol): - def __call__(self, document: Mapping[str, object]) -> dict[str, str]: ... - - -class NativeUploadDocument(Protocol): - def __call__(self, file_content: bytes, file_name: str | None, content_type: str | None) -> dict[str, str]: ... - - -class NativeMimeType(Protocol): - def __call__(self, file_name: str) -> str: ... - - -_FILE_DOCUMENT: Final = NativeBinding( - "_ocr_file_document", - validate=lambda value: ( - cast( # cast-ok: native export owns the callable signature - NativeFileDocument, value - ) - if callable(value) - else None - ), -) -_UPLOAD_DOCUMENT: Final = NativeBinding( - "_ocr_upload_document", - validate=lambda value: ( - cast( # cast-ok: native export owns the callable signature - NativeUploadDocument, value - ) - if callable(value) - else None - ), -) -_MAX_FILE_BYTES: Final = NativeBinding( - "_OCR_MAX_FILE_BYTES", validate=lambda value: value if isinstance(value, int) and value > 0 else None -) -_MIME_TYPE: Final = NativeBinding( - "_ocr_mime_type", - validate=lambda value: ( - cast( # cast-ok: native export owns the callable signature - NativeMimeType, value - ) - if callable(value) - else None - ), -) -_PYTHON_MAX_FILE_BYTES: Final = 50 * 1024 * 1024 - - -def get_mime_type(file_path: str) -> str: - native: Final = _MIME_TYPE.load() if rust_ocr_enabled() else None - if native is None: - from litellm.ocr import legacy - - return legacy.get_mime_type(file_path) - return native(file_path) - - -def get_max_file_bytes() -> int: - limit: Final = _MAX_FILE_BYTES.load() if rust_ocr_enabled() else None - if limit is None: - return _PYTHON_MAX_FILE_BYTES - return limit - - -def convert_file_document_to_url_document(document: FileDocument) -> dict[str, str]: - native: Final = _FILE_DOCUMENT.load() if rust_ocr_enabled() else None - if native is None: - from litellm.ocr import legacy - - return legacy.convert_file_document_to_url_document(document) - return native(document) - - -def convert_upload_to_url_document( - file_content: bytes, filename: str | None, content_type: str | None -) -> dict[str, str]: - native: Final = _UPLOAD_DOCUMENT.load() if rust_ocr_enabled() else None - if native is None: - from litellm.ocr import legacy - - if len(file_content) > _PYTHON_MAX_FILE_BYTES: - raise ValueError("OCR file exceeds the size limit") - content_mime: Final = content_type.split(";")[0].strip() if content_type else None - mime_type: Final = ( - legacy.get_mime_type(filename) - if filename and (not content_mime or content_mime == "application/octet-stream") - else content_mime or "application/octet-stream" - ) - return legacy.convert_file_document_to_url_document( - {"type": "file", "file": file_content, "mime_type": mime_type} - ) - return native(file_content, filename, content_type) diff --git a/litellm/ocr/legacy.py b/litellm/ocr/legacy.py index a742be274b3..f0cf6cc82cc 100644 --- a/litellm/ocr/legacy.py +++ b/litellm/ocr/legacy.py @@ -11,7 +11,7 @@ from collections.abc import Coroutine, Mapping from dataclasses import dataclass from io import IOBase from types import MappingProxyType -from typing import Final, cast # noqa: TID251 # adapters preserve the legacy untyped contracts +from typing import Final, Protocol, cast # noqa: TID251 # adapters preserve the legacy untyped contracts import httpx @@ -26,7 +26,6 @@ from litellm.llms.base_llm.ocr.transformation import ( parse_ocr_request_format, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.ocr.input import FileReader from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CustomPricingLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -34,6 +33,10 @@ from litellm.utils import ProviderConfigManager, client base_llm_http_handler: Final = BaseLLMHTTPHandler() +class FileReader(Protocol): + def read(self) -> bytes | str: ... + + @dataclass(frozen=True, slots=True) class _PreparedOCRRequest: model: str diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 382c5d6aae4..c6371c0c33f 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -5,7 +5,7 @@ import httpx from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.ocr import legacy -from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type +from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type from litellm.rust_bridge.bindings import native_exception_types from litellm.rust_bridge.configuration import rust_ocr_enabled from litellm.rust_bridge.ocr import LiteLLMOcrRequest diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index 53ebbe91b54..dde3d5ceb50 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -15,12 +15,13 @@ from litellm.llms.base_llm.ocr.transformation import ( OCRResponse, parse_ocr_request_format, ) -from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes +from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing router: Final = APIRouter() +_MAX_FILE_BYTES: Final = 50 * 1024 * 1024 def _build_document_from_upload( @@ -28,7 +29,15 @@ def _build_document_from_upload( filename: str | None, content_type: str | None, ) -> dict[str, str]: - return convert_upload_to_url_document(file_content, filename, content_type) + supplied_mime: Final = content_type.split(";")[0].strip() if content_type else None + mime_type: Final = ( + get_mime_type(filename) + if filename and (not supplied_mime or supplied_mime == "application/octet-stream") + else supplied_mime + ) + return convert_file_document_to_url_document( + {"type": "file", "file": file_content, "mime_type": mime_type or "application/octet-stream"} + ) def _with_request_format(data: Mapping[str, Any], request: Request) -> Mapping[str, Any]: @@ -103,9 +112,11 @@ async def _parse_multipart_form(request: Request) -> dict[str, Any]: # Seek to start in case the file was already partially read by middleware await uploaded_file.seek(0) - file_content: Final = await uploaded_file.read(get_max_file_bytes() + 1) + file_content: Final = await uploaded_file.read(_MAX_FILE_BYTES + 1) if not file_content: raise ValueError("Uploaded file is empty") + if len(file_content) > _MAX_FILE_BYTES: + raise ValueError("OCR file exceeds the size limit") document: Final = _build_document_from_upload( file_content=file_content, diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index e62c85f4599..32b20bb7931 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -33,15 +33,6 @@ def aocr( timeout_seconds: float | None = None, ) -> Future[dict[str, object]]: ... -_OCR_MAX_FILE_BYTES: int - -def _ocr_upload_document( - file_content: bytes, - file_name: str | None = None, - content_type: str | None = None, -) -> dict[str, str]: ... -def _ocr_file_document(document: Mapping[str, object]) -> dict[str, str]: ... -def _ocr_mime_type(file_name: str) -> str: ... def _ocr_lifecycle( request: LiteLLMOcrRequest, args: tuple[object, ...], @@ -139,15 +130,11 @@ class TokenCounter: def gil_stats() -> dict[str, int]: ... __all__ = [ - "_OCR_MAX_FILE_BYTES", "ResponsesWebSocketConnection", "RustBridgeDeclined", "RustUpstreamError", "TokenCounter", - "_ocr_file_document", "_ocr_lifecycle", - "_ocr_mime_type", - "_ocr_upload_document", "achat_completions", "amessages", "aocr", diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/test_litellm/ocr/test_ocr_file_input.py index 3526d8c00d6..8f82a64bd85 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/tests/test_litellm/ocr/test_ocr_file_input.py @@ -12,32 +12,16 @@ Tests that: import base64 import os import tempfile -from collections.abc import Generator from io import BytesIO from pathlib import Path from typing import Final -from unittest.mock import AsyncMock, MagicMock, Mock +from unittest.mock import AsyncMock, MagicMock import orjson import pytest from starlette.datastructures import FormData -from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type - - -@pytest.fixture(autouse=True, params=["native", "disabled", "unavailable"]) -def document_runtime(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Generator[None]: - from litellm.rust_bridge import bindings, configuration - - configuration.reset_rust_configuration() - monkeypatch.delenv("LITELLM_RUST", raising=False) - if request.param == "disabled": - monkeypatch.setenv("LITELLM_RUST", "0") - monkeypatch.setattr(bindings, "get_native_bridge", Mock(side_effect=AssertionError("Rust is disabled"))) - elif request.param == "unavailable": - monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) - yield - configuration.reset_rust_configuration() +from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type class TestGetMimeType: @@ -503,10 +487,9 @@ class TestProxySecurityGuard: async def test_proxy_upload_stops_reading_at_size_limit() -> None: from starlette.datastructures import UploadFile - from litellm.ocr.input import get_max_file_bytes - from litellm.proxy.ocr_endpoints.endpoints import _parse_multipart_form + from litellm.proxy.ocr_endpoints.endpoints import _MAX_FILE_BYTES, _parse_multipart_form - limit: Final = get_max_file_bytes() + limit: Final = _MAX_FILE_BYTES with tempfile.TemporaryFile() as stream: stream.truncate(limit * 2) upload: Final = UploadFile(file=stream, filename="large.pdf") diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 4f4b39fa6c6..58bb6a77537 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -518,32 +518,34 @@ async def test_native_ocr_inherits_named_credentials_without_overwriting_argumen assert ocr_server.requests[0].body["pages"] == [0, 2] -@pytest.mark.parametrize("source", ["sdk", "proxy"]) @pytest.mark.parametrize( - "filename,mime", [("scan.PNG", "image/png"), ("document.pdf", "application/pdf"), ("note.txt", "text/plain")] + "filename,field,mime", + [("scan.PNG", "image_url", "image/png"), ("document.pdf", "document_url", "application/pdf")], ) -def test_ocr_file_helpers_use_native_document_preparation(source: str, filename: str, mime: str) -> None: +def test_native_ocr_infers_mime_type_from_reader_name( + ocr_server: RecordingServer, filename: str, field: str, mime: str +) -> None: from io import BytesIO - from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type - from litellm.proxy.ocr_endpoints.endpoints import _build_document_from_upload - file: Final = BytesIO(b"abc") file.name = filename - document: Final = ( - convert_file_document_to_url_document({"type": "file", "file": file}) - if source == "sdk" - else _build_document_from_upload(b"abc", filename, "application/octet-stream; charset=utf-8") - ) - field: Final = "image_url" if mime.startswith("image/") else "document_url" - assert get_mime_type(filename) == mime - assert document == {"type": field, field: f"data:{mime};base64,YWJj"} + call_native_ocr(ocr_server, document={"type": "file", "file": file}) + assert ocr_server.requests[0].body["document"] == {"type": field, field: f"data:{mime};base64,YWJj"} + + +def test_native_ocr_encodes_str_reader_results_as_utf8(ocr_server: RecordingServer) -> None: + from io import StringIO + + call_native_ocr(ocr_server, document={"type": "file", "file": StringIO("abc"), "mime_type": "text/plain"}) + assert ocr_server.requests[0].body["document"] == { + "type": "document_url", + "document_url": "data:text/plain;base64,YWJj", + } @pytest.mark.parametrize("attribute", ["read", "name"]) -def test_native_file_preparation_preserves_property_errors(attribute: str) -> None: - from litellm.ocr.input import convert_file_document_to_url_document - +def test_native_file_preparation_preserves_property_errors(ocr_server: RecordingServer, attribute: str) -> None: + ocr_server.expected_requests = 0 failure: Final = LookupError("file property failed") class File: @@ -555,16 +557,47 @@ def test_native_file_preparation_preserves_property_errors(attribute: str) -> No def read(self): return b"abc" - with pytest.raises(LookupError) as caught: - convert_file_document_to_url_document({"type": "file", "file": File()}) - assert caught.value is failure + with pytest.raises(litellm.APIConnectionError, match="file property failed") as caught: + call_native_ocr(ocr_server, document={"type": "file", "file": File()}) + assert caught.value.__context__ is failure + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_native_file_preparation_preserves_reader_exception( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + ocr_server.expected_requests = 0 + failure: Final = RuntimeError("reader failed") + + class Reader: + def read(self) -> bytes: + raise failure + + document: Final = {"type": "file", "file": Reader()} + with pytest.raises(litellm.APIConnectionError, match="reader failed") as caught: + await call_native_aocr(ocr_server, document=document) if asynchronous else call_native_ocr( + ocr_server, document=document + ) + assert caught.value.__context__ is failure + + +def test_native_file_preparation_rejects_unsupported_reader_results(ocr_server: RecordingServer) -> None: + ocr_server.expected_requests = 0 + + class Reader: + def read(self) -> int: + return 1 + + with pytest.raises(litellm.APIConnectionError, match="bytes or str") as caught: + call_native_ocr(ocr_server, document={"type": "file", "file": Reader()}) + assert isinstance(caught.value.__context__, TypeError) @pytest.mark.parametrize("kind", ["bytes", "path", "reader"]) -def test_native_file_preparation_rejects_oversized_input(kind: str, tmp_path: Path) -> None: - from litellm.ocr.input import FileDocument, convert_file_document_to_url_document, get_max_file_bytes - - limit: Final = get_max_file_bytes() +def test_native_file_preparation_rejects_oversized_input(ocr_server: RecordingServer, kind: str, tmp_path: Path) -> None: + ocr_server.expected_requests = 0 + limit: Final = 50 * 1024 * 1024 path: Final = tmp_path / "large.pdf" with path.open("wb") as stream: stream.truncate(limit + 1) @@ -573,53 +606,25 @@ def test_native_file_preparation_rejects_oversized_input(kind: str, tmp_path: Pa def read(self) -> bytes: return b"a" * (limit + 1) - document: Final[FileDocument] = { + document: Final = { "type": "file", "file": path if kind == "path" else Reader() if kind == "reader" else b"a" * (limit + 1), } - with pytest.raises(ValueError, match="exceeds the size limit"): - convert_file_document_to_url_document(document) + with pytest.raises(litellm.BadRequestError, match="exceeds the size limit"): + call_native_ocr(ocr_server, document=document) -@pytest.mark.parametrize("kind", ["str", "path", "reader"]) -def test_native_upload_binding_rejects_filesystem_inputs(kind: str, tmp_path: Path) -> None: +def test_native_file_preparation_reports_missing_paths(ocr_server: RecordingServer, tmp_path: Path) -> None: + ocr_server.expected_requests = 0 + missing: Final = tmp_path / "missing.pdf" + with pytest.raises(litellm.APIConnectionError, match=f"File not found: {missing}") as caught: + call_native_ocr(ocr_server, document={"type": "file", "file": missing}) + assert isinstance(caught.value.__context__, FileNotFoundError) + + +def test_native_file_preparation_rejects_empty_readers(ocr_server: RecordingServer) -> None: from io import BytesIO - from typing import cast # noqa: TID251 # deliberately invalid inputs exercise the native runtime boundary - from litellm.ocr.input import convert_upload_to_url_document - - path: Final = tmp_path / "secret.pdf" - path.write_bytes(b"server secret") - source: Final = str(path) if kind == "str" else path if kind == "path" else BytesIO(b"abc") - with pytest.raises(TypeError): - convert_upload_to_url_document(cast(bytes, source), "document.pdf", None) - - -@pytest.mark.parametrize("extra_bytes", [0, 1]) -def test_native_upload_enforces_file_size_limit(extra_bytes: int) -> None: - import base64 - - from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes - - content: Final = b"a" * (get_max_file_bytes() + extra_bytes) - if extra_bytes: - with pytest.raises(ValueError, match="exceeds the size limit"): - convert_upload_to_url_document(content, "scan.pdf", None) - return - document: Final = convert_upload_to_url_document(content, "scan.pdf", None) - assert document["type"] == "document_url" - assert base64.b64decode(document["document_url"].split(",", 1)[1]) == content - - -def test_native_file_preparation_preserves_reader_exception() -> None: - from litellm.ocr.input import convert_file_document_to_url_document - - failure: Final = RuntimeError("reader failed") - - class Reader: - def read(self) -> bytes: - raise failure - - with pytest.raises(RuntimeError) as caught: - convert_file_document_to_url_document({"type": "file", "file": Reader()}) - assert caught.value is failure + ocr_server.expected_requests = 0 + with pytest.raises(litellm.BadRequestError, match="File is empty"): + call_native_ocr(ocr_server, document={"type": "file", "file": BytesIO(b"")}) From cfe65f7b551c192fb6edd402ea37ba7ea4646e18 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 14:15:27 -0700 Subject: [PATCH 128/187] fix(proxy)!: enforce fallback budget by default A budget bypass that ships off by default stays open for every deployment that does not know to look for the flag, so `enforce_fallback_budget` now defaults to true and `general_settings.enforce_fallback_budget: false` is the opt-out for anyone who wants the old unguarded behaviour back. BREAKING CHANGE: a paid fallback target is now refused for callers who are over their key or user `max_budget`. Deployments relying on fallbacks to keep serving over-budget callers must set enforce_fallback_budget: false. --- litellm/proxy/auth/fallback_budget.py | 6 +++--- .../proxy/auth/test_fallback_budget.py | 18 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 9 +++++---- 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/fallback_budget.py b/litellm/proxy/auth/fallback_budget.py index 00e885d8d88..e356f8acc7d 100644 --- a/litellm/proxy/auth/fallback_budget.py +++ b/litellm/proxy/auth/fallback_budget.py @@ -8,8 +8,8 @@ actually bills. So a free model with a paid fallback spends without a gate. This predicate is injected into the router to re-check budget for each fallback target before it is attempted, mirroring `fallback_model_access.py`. It deliberately leaves the primary attempt alone: -a zero-cost model is never blocked by budget, and only the paid fallback is refused. Opt-in via -`general_settings.enforce_fallback_budget: true`. +a zero-cost model is never blocked by budget, and only the paid fallback is refused. On by default; +set `general_settings.enforce_fallback_budget: false` to restore the unguarded behaviour. Scope: the key's and the user's `max_budget`. Not covered yet, and each needs a read-only evaluation path before it can be: team, team-member, end-user, org, global and per-model budgets, whose @@ -50,7 +50,7 @@ class _RequestMetadata(BaseModel): class _FallbackBudgetSettings(BaseModel): - enforce_fallback_budget: bool = False + enforce_fallback_budget: bool = True def _token_in_metadata(metadata: object) -> UserAPIKeyAuth | None: diff --git a/tests/test_litellm/proxy/auth/test_fallback_budget.py b/tests/test_litellm/proxy/auth/test_fallback_budget.py index 0ff1e05826d..00c1a7cdefc 100644 --- a/tests/test_litellm/proxy/auth/test_fallback_budget.py +++ b/tests/test_litellm/proxy/auth/test_fallback_budget.py @@ -5,6 +5,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.fallback_budget import ( RouterFallbackBudgetCheck, is_token_within_budget_for_model, + router_fallback_budget_check, ) FREE_MODEL = { @@ -182,3 +183,20 @@ async def test_router_without_a_budget_check_attempts_every_fallback(): over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", over) is True + + +@pytest.mark.asyncio +async def test_enforcement_is_on_by_default_and_opt_out_restores_the_leak(monkeypatch): + """ + Leaving the paid fallback unguarded is the budget bypass this module exists to close, so an + unconfigured proxy has to enforce. `enforce_fallback_budget: false` is the deliberate opt-out. + """ + from litellm.proxy import proxy_server + + over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + + monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) + assert await router_fallback_budget_check(model="paid-model", request_kwargs=over, llm_router=_router()) is False + + monkeypatch.setattr(proxy_server, "general_settings", {"enforce_fallback_budget": False}, raising=False) + assert await router_fallback_budget_check(model="paid-model", request_kwargs=over, llm_router=_router()) is True diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index ca53887623f..ad28b6b4421 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13416,14 +13416,15 @@ async def test_load_config_router_budget_checks_fallback_targets_against_the_cal } } - # off by default: the paid fallback is still attempted for an over-budget caller + # on by default: an over-budget caller is refused the paid fallback with no config at all monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) - assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, llm_router=router) is True - - monkeypatch.setattr(proxy_server, "general_settings", {"enforce_fallback_budget": True}, raising=False) assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, llm_router=router) is False assert await router.fallback_budget_check(model="m", request_kwargs=under_budget, llm_router=router) is True + # explicit opt-out restores the unguarded behaviour + monkeypatch.setattr(proxy_server, "general_settings", {"enforce_fallback_budget": False}, raising=False) + assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, llm_router=router) is True + @pytest.mark.asyncio async def test_load_config_user_api_key_cache_max_size_keeps_more_than_200_entries(tmp_path, monkeypatch): From b4781317012c57c1bd7e186d1478b465316c6dfb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 14:20:49 -0700 Subject: [PATCH 129/187] test(together_ai): select the live model from the cost map `together_ai/openai/gpt-oss-20b` was hardcoded in two live tests and is no longer served, so both failed on a vendor catalog change rather than on anything litellm did. Both call sites now resolve the cheapest non-deprecated together_ai chat entry at runtime, filtered on the capabilities the tests actually exercise, mirroring what tests/e2e/llm_translation/test_together_ai_e2e.py already does. The selector lives in tests/_live_test_helpers.py so both lanes share one implementation. --- tests/_live_test_helpers.py | 36 +++++++++++++++++++++ tests/llm_translation/test_together_ai.py | 7 +++- tests/local_testing/test_text_completion.py | 7 +++- 3 files changed, 48 insertions(+), 2 deletions(-) diff --git a/tests/_live_test_helpers.py b/tests/_live_test_helpers.py index a79b81e82c1..6b39f37921d 100644 --- a/tests/_live_test_helpers.py +++ b/tests/_live_test_helpers.py @@ -1,4 +1,7 @@ import os +from collections.abc import Mapping +from datetime import date +from typing import Any import pytest @@ -8,3 +11,36 @@ def _skip_live_prompt_caching_test(): pytest.skip("Live prompt-caching E2E tests are opt-in") if os.environ.get("CASSETTE_REDIS_URL"): pytest.skip("Live prompt-caching E2E tests cannot run under VCR replay") + + +def cheapest_together_chat_model(*capability_flags: str) -> str: + import litellm + + today = date.today().isoformat() + + def qualifies(name: str, entry: Mapping[str, Any]) -> bool: + deprecation_date = entry.get("deprecation_date") + return ( + name.startswith("together_ai/") + and entry.get("litellm_provider") == "together_ai" + and entry.get("mode") == "chat" + and (deprecation_date is None or deprecation_date > today) + and (entry.get("input_cost_per_token") or 0.0) > 0 + and (entry.get("output_cost_per_token") or 0.0) > 0 + and all(bool(entry.get(flag)) for flag in capability_flags) + ) + + candidates = sorted( + ( + name + for name, entry in litellm.model_cost.items() + if isinstance(entry, Mapping) and qualifies(name, entry) + ), + key=lambda name: ( + litellm.model_cost[name].get("input_cost_per_token") or 0.0, + litellm.model_cost[name].get("output_cost_per_token") or 0.0, + name, + ), + ) + assert candidates, f"no live together_ai chat model in the cost map satisfies {capability_flags}" + return candidates[0] diff --git a/tests/llm_translation/test_together_ai.py b/tests/llm_translation/test_together_ai.py index fd7ad40ed11..7a49d46b528 100644 --- a/tests/llm_translation/test_together_ai.py +++ b/tests/llm_translation/test_together_ai.py @@ -3,6 +3,7 @@ Test TogetherAI LLM """ from base_llm_unit_tests import BaseLLMChatTest +from tests._live_test_helpers import cheapest_together_chat_model import json import os from datetime import datetime @@ -16,7 +17,11 @@ import pytest class TestTogetherAI(BaseLLMChatTest): def get_base_completion_call_args(self) -> dict: litellm.set_verbose = True - return {"model": "together_ai/openai/gpt-oss-20b"} + return { + "model": cheapest_together_chat_model( + "supports_function_calling", "supports_response_schema" + ) + } def test_tool_call_no_arguments(self, tool_call_no_arguments): """Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833""" diff --git a/tests/local_testing/test_text_completion.py b/tests/local_testing/test_text_completion.py index a814ce6d303..b15037a2fcd 100644 --- a/tests/local_testing/test_text_completion.py +++ b/tests/local_testing/test_text_completion.py @@ -1,7 +1,11 @@ import asyncio import json +import os +import sys import traceback +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) + from dotenv import load_dotenv load_dotenv() @@ -12,6 +16,7 @@ from unittest.mock import MagicMock, patch import pytest import litellm +from tests._live_test_helpers import cheapest_together_chat_model from litellm import ( RateLimitError, TextCompletionResponse, @@ -4030,7 +4035,7 @@ def test_async_text_completion_together_ai(): async def test_get_response(): try: response = await litellm.atext_completion( - model="together_ai/openai/gpt-oss-20b", + model=cheapest_together_chat_model(), prompt="good morning", max_tokens=10, ) From 09f3a5160a233556876460718607b440cc35aa37 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:21:14 -0700 Subject: [PATCH 130/187] test(azure_ai): assert the bare deployment name reaches the native Responses endpoint --- .../test_azure_ai_responses_transformation.py | 86 ++++++++++++++----- 1 file changed, 66 insertions(+), 20 deletions(-) diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py index 1fe604d4a6b..925608a3c9b 100644 --- a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py +++ b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -192,21 +192,54 @@ def test_validate_environment_raises_without_credentials(): ) +NATIVE_RESPONSES_CASES = [ + ("azure_ai/gpt-5.6-luna-20260710154139", FOUNDRY_PROJECT_BASE, FOUNDRY_RESPONSES_URL, "gpt-5.6-luna-20260710154139"), + ( + "azure_ai/gpt-5.6-luna", + "https://res.services.ai.azure.com/models", + "https://res.services.ai.azure.com/openai/v1/responses", + "gpt-5.6-luna", + ), + ( + "azure_ai/gpt-5.6-sol", + "https://res.services.ai.azure.com", + "https://res.services.ai.azure.com/openai/v1/responses", + "gpt-5.6-sol", + ), + ( + "azure_ai/gpt-5.6-luna-20260710154139", + "https://res.openai.azure.com", + "https://res.openai.azure.com/openai/v1/responses", + "gpt-5.6-luna-20260710154139", + ), + ( + "azure_ai/gpt-5.6-sol", + "https://res.openai.azure.com", + "https://res.openai.azure.com/openai/v1/responses", + "gpt-5.6-sol", + ), +] + + +def _assert_native_responses_request(route, expected_url, expected_model): + request = route.calls.last.request + body = json.loads(request.content) + assert f"{request.url.scheme}://{request.url.host}{request.url.path}" == expected_url + assert request.headers["api-key"] == "fake-key" + assert body["model"] == expected_model + assert body["input"] == "What is the weather in SF?" + assert "messages" not in body + assert body["reasoning"] == {"effort": "high"} + assert body["tools"] == [WEATHER_TOOL] + + @pytest.mark.asyncio @respx.mock -@pytest.mark.parametrize( - "model,api_base,expected_url", - [ - ("azure_ai/gpt-5.6-luna-20260710154139", FOUNDRY_PROJECT_BASE, FOUNDRY_RESPONSES_URL), - ( - "azure_ai/gpt-5.6-luna", - "https://res.services.ai.azure.com/models", - "https://res.services.ai.azure.com/openai/v1/responses", - ), - ], -) -async def test_aresponses_sends_reasoning_and_tools_to_native_endpoint(model, api_base, expected_url): - route = respx.post(expected_url).mock(return_value=httpx.Response(200, json=_responses_payload("gpt-5.6-luna"))) +@pytest.mark.parametrize("model,api_base,expected_url,expected_model", NATIVE_RESPONSES_CASES) +async def test_aresponses_sends_reasoning_and_tools_to_native_endpoint(model, api_base, expected_url, expected_model): + route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock( + return_value=httpx.Response(200, json=_responses_payload(expected_model)) + ) await litellm.aresponses( model=model, @@ -217,13 +250,26 @@ async def test_aresponses_sends_reasoning_and_tools_to_native_endpoint(model, ap api_key="fake-key", ) - request = route.calls.last.request - body = json.loads(request.content) - assert request.headers["api-key"] == "fake-key" - assert body["input"] == "What is the weather in SF?" - assert "messages" not in body - assert body["reasoning"] == {"effort": "high"} - assert body["tools"] == [WEATHER_TOOL] + _assert_native_responses_request(route, expected_url, expected_model) + + +@pytest.mark.asyncio +@respx.mock +@pytest.mark.parametrize("model,api_base,expected_url,expected_model", NATIVE_RESPONSES_CASES) +async def test_router_aresponses_sends_bare_deployment_name(model, api_base, expected_url, expected_model): + route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock( + return_value=httpx.Response(200, json=_responses_payload(expected_model)) + ) + router = litellm.Router( + model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": model, "api_base": api_base, "api_key": "fake-key"}}], + num_retries=0, + ) + + await router.aresponses( + model="gpt-5.6", input="What is the weather in SF?", reasoning={"effort": "high"}, tools=[WEATHER_TOOL] + ) + + _assert_native_responses_request(route, expected_url, expected_model) @pytest.mark.asyncio From ba6c9fa61d87bd76634bdaba7ce4ef578771f923 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 14:26:16 -0700 Subject: [PATCH 131/187] fix(test): drop the redundant sys.path.insert CI runs these lanes as `python -m pytest` from the repo root, so the root is already on sys.path and `tests._live_test_helpers` imports without help. The insert only tripped the TQ003 test-quality budget. --- tests/local_testing/test_text_completion.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/local_testing/test_text_completion.py b/tests/local_testing/test_text_completion.py index b15037a2fcd..6808dfd768b 100644 --- a/tests/local_testing/test_text_completion.py +++ b/tests/local_testing/test_text_completion.py @@ -1,11 +1,7 @@ import asyncio import json -import os -import sys import traceback -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) - from dotenv import load_dotenv load_dotenv() From 515bf8c9d564731cfb285d47818db2fc081e1ba7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 14:32:08 -0700 Subject: [PATCH 132/187] refactor(test): validate cost-map entries into a typed model The selector read raw cost-map dicts as `Mapping[str, Any]`. It now validates each together_ai entry into a frozen Pydantic model and takes the two capabilities as keyword booleans, so nothing in the helper is coarsely typed or stringly addressed. --- tests/_live_test_helpers.py | 56 +++++++++++++++-------- tests/llm_translation/test_together_ai.py | 2 +- 2 files changed, 38 insertions(+), 20 deletions(-) diff --git a/tests/_live_test_helpers.py b/tests/_live_test_helpers.py index 6b39f37921d..629f8ac9fdb 100644 --- a/tests/_live_test_helpers.py +++ b/tests/_live_test_helpers.py @@ -1,9 +1,8 @@ import os -from collections.abc import Mapping from datetime import date -from typing import Any import pytest +from pydantic import BaseModel, ConfigDict def _skip_live_prompt_caching_test(): @@ -13,34 +12,53 @@ def _skip_live_prompt_caching_test(): pytest.skip("Live prompt-caching E2E tests cannot run under VCR replay") -def cheapest_together_chat_model(*capability_flags: str) -> str: + +class TogetherCostEntry(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + litellm_provider: str | None = None + mode: str | None = None + deprecation_date: str | None = None + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + supports_function_calling: bool | None = None + supports_response_schema: bool | None = None + + +def cheapest_together_chat_model( + *, function_calling: bool = False, response_schema: bool = False +) -> str: import litellm today = date.today().isoformat() - def qualifies(name: str, entry: Mapping[str, Any]) -> bool: - deprecation_date = entry.get("deprecation_date") + def qualifies(name: str, entry: TogetherCostEntry) -> bool: return ( name.startswith("together_ai/") - and entry.get("litellm_provider") == "together_ai" - and entry.get("mode") == "chat" - and (deprecation_date is None or deprecation_date > today) - and (entry.get("input_cost_per_token") or 0.0) > 0 - and (entry.get("output_cost_per_token") or 0.0) > 0 - and all(bool(entry.get(flag)) for flag in capability_flags) + and entry.litellm_provider == "together_ai" + and entry.mode == "chat" + and (entry.deprecation_date is None or entry.deprecation_date > today) + and (entry.input_cost_per_token or 0.0) > 0 + and (entry.output_cost_per_token or 0.0) > 0 + and (not function_calling or bool(entry.supports_function_calling)) + and (not response_schema or bool(entry.supports_response_schema)) ) + registry: dict[str, TogetherCostEntry] = { + name: TogetherCostEntry.model_validate(raw) + for name, raw in litellm.model_cost.items() + if isinstance(raw, dict) and name.startswith("together_ai/") + } candidates = sorted( - ( - name - for name, entry in litellm.model_cost.items() - if isinstance(entry, Mapping) and qualifies(name, entry) - ), + (name for name, entry in registry.items() if qualifies(name, entry)), key=lambda name: ( - litellm.model_cost[name].get("input_cost_per_token") or 0.0, - litellm.model_cost[name].get("output_cost_per_token") or 0.0, + registry[name].input_cost_per_token or 0.0, + registry[name].output_cost_per_token or 0.0, name, ), ) - assert candidates, f"no live together_ai chat model in the cost map satisfies {capability_flags}" + assert candidates, ( + "no live together_ai chat model in the cost map satisfies " + f"function_calling={function_calling} response_schema={response_schema}" + ) return candidates[0] diff --git a/tests/llm_translation/test_together_ai.py b/tests/llm_translation/test_together_ai.py index 7a49d46b528..0b4e9d3952c 100644 --- a/tests/llm_translation/test_together_ai.py +++ b/tests/llm_translation/test_together_ai.py @@ -19,7 +19,7 @@ class TestTogetherAI(BaseLLMChatTest): litellm.set_verbose = True return { "model": cheapest_together_chat_model( - "supports_function_calling", "supports_response_schema" + function_calling=True, response_schema=True ) } From 0de187e76cc0268bdf8e19e73b3cdaac80015bd2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 14:34:13 -0700 Subject: [PATCH 133/187] style(test): annotate the new locals as Final --- .../test_generic_api_callback.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/logging_callback_tests/test_generic_api_callback.py b/tests/logging_callback_tests/test_generic_api_callback.py index 2c62741ed38..d9853ebcb52 100644 --- a/tests/logging_callback_tests/test_generic_api_callback.py +++ b/tests/logging_callback_tests/test_generic_api_callback.py @@ -10,6 +10,7 @@ import httpx import json import logging import time +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -98,8 +99,8 @@ async def test_generic_api_callback(): assert isinstance(actual_request, list), "Request body should be a list" assert len(actual_request) > 0, "Request body list should not be empty" - this_test_messages = [{"role": "user", "content": "Hello, world!"}] - mine = [ + this_test_messages: Final = [{"role": "user", "content": "Hello, world!"}] + mine: Final = [ item for item in actual_request if item.get("messages") == this_test_messages ] assert ( @@ -455,12 +456,14 @@ async def test_generic_api_callback_sumologic_uses_ndjson(): assert isinstance(ndjson_data, str), "Data should be a string for NDJSON" lines = ndjson_data.strip().split("\n") - records = [json.loads(line) for line in lines] + records: Final = [json.loads(line) for line in lines] - this_test_messages = [ + this_test_messages: Final = [ [{"role": "user", "content": f"Test {i}"}] for i in range(2) ] - mine = [record for record in records if record.get("messages") in this_test_messages] + mine: Final = [ + record for record in records if record.get("messages") in this_test_messages + ] assert ( len(mine) == 2 ), f"Expected this test's 2 calls as NDJSON lines, got {len(mine)} of {len(records)}" From a6b10ad654948fd70d84c697542703cd3b4e2a0f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 14:58:57 -0700 Subject: [PATCH 134/187] fix(prices): dedupe Nova cache_read_input_token_cost keys left by a text merge PR #41343 and PR #41112 both added cache_read_input_token_cost to the six amazon.nova-{micro,lite,pro}-v1:0 and us.amazon.nova-* entries, one at the top of each entry and one at the bottom. The merge kept both, so every PR now fails test_price_map_has_no_duplicate_keys. Both copies carried the same value, so this only removes the trailing duplicate in both price files --- ...model_prices_and_context_window_backup.json | 18 ++++++------------ model_prices_and_context_window.json | 18 ++++++------------ 2 files changed, 12 insertions(+), 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b1c38d0350a..92e5b1c4ff7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -377,8 +377,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -561,8 +560,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -578,8 +576,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45794,8 +45791,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45809,8 +45805,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45842,8 +45837,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b1c38d0350a..92e5b1c4ff7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -377,8 +377,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -561,8 +560,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -578,8 +576,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45794,8 +45791,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45809,8 +45805,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45842,8 +45837,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, From 76c0f8db1d415a0307f4188a0ca992688ec3b44b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 15:01:08 -0700 Subject: [PATCH 135/187] chore(e2e): report the key components behind a mount that never converges Builds 232 and 233 held the Bedrock hit rate at 9% with the Claude Code driver already sending byte-identical requests and headers, so something between the proxy's ingress and the upstream still moves per build and the flat key cannot say what. Emit a digest per key component next to the counters: the test id, the method, the URL, each keyed header, the whole body, and one digest per top-level JSON body field. Values are digested, so no payload or credential reaches the artifact. Diffing two builds' artifacts names the field that moved. Diagnostic, to be removed once it has answered. --- tests/e2e/provider_cache.py | 63 +++++++++++++++++++++++++++++-- tests/e2e/provider_cache_redis.py | 4 ++ 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/tests/e2e/provider_cache.py b/tests/e2e/provider_cache.py index 444972ffce5..22967869c78 100644 --- a/tests/e2e/provider_cache.py +++ b/tests/e2e/provider_cache.py @@ -4,6 +4,8 @@ import base64 import hashlib import hmac import io +import json +import os import threading import time from collections.abc import Callable, Generator, Mapping @@ -400,6 +402,51 @@ def decode_response(secret: bytes, key: str, payload: bytes, mount: str, url: st return response +def component_digests( + test_key: str, method: str, url: str, headers: Mapping[str, str], body: bytes | None, +) -> dict[str, str]: + """Per-component digests of everything the key covers. + + A mount whose corpus never converges is a mount where one of these moves + between builds, and the flat key cannot say which. Values are digested, so + no payload or credential is written, and a JSON body contributes one digest + per top-level field so the field that moved can be named.""" + parts: dict[str, str] = { # rebind-ok: a report assembled from three differently shaped sources + "test_key": test_key, + "method": method, + "url": short_digest(canonical_text(url).encode()), + } + for name, value in sorted(headers.items()): + parts[f"header:{name.lower()}"] = short_digest(value.encode()) + canonical: Final = b"" if body is None else canonical_body(body) + parts["body"] = short_digest(canonical) + try: + parsed: Final = JSON_VALUE.validate_json(canonical) + except ValidationError: + return parts + if isinstance(parsed, dict): + for name, value in sorted(parsed.items()): + parts[f"body:{name}"] = short_digest(json.dumps(value, sort_keys=True).encode()) + return parts + + +def short_digest(value: bytes) -> str: + return hashlib.sha256(value).hexdigest()[:16] + + +@dataclass(slots=True) +class KeyProbe: + """Every keyed request's components, when a metrics directory is configured.""" + + rows: tuple[tuple[tuple[str, str], ...], ...] = () + lock: threading.Lock = field(default_factory=threading.Lock) + + def observe(self, mount: str, outcome: str, parts: Mapping[str, str]) -> None: + row: Final = tuple({"mount": mount, "outcome": outcome, **parts}.items()) + with self.lock: + self.rows = (*self.rows, row) + + @dataclass(slots=True) class CacheCounters: counts: tuple[tuple[str, int], ...] = () @@ -464,6 +511,7 @@ class CacheEdge: store: ResponseStore secret: bytes = field(repr=False) counters: CacheCounters = field(default_factory=CacheCounters) + probe: KeyProbe = field(default_factory=KeyProbe) slots: SlotCounter = field(default_factory=SlotCounter) policies: Mapping[str, MountPolicy] = NO_POLICIES wait_seconds: float = 2.0 @@ -481,6 +529,14 @@ class CacheEdge: self.counters.increment(name) self.counters.increment(f"mount:{mount}:{name}") + def record_key( + self, mount: str, outcome: str, test_key: str, method: str, url: str, + headers: Mapping[str, str], body: bytes | None, + ) -> None: + if not os.environ.get("E2E_PROVIDER_CACHE_METRICS_DIR"): + return + self.probe.observe(mount, outcome, component_digests(test_key, method, url, headers, body)) + def outbound(self, mount: str, method: str, url: str, headers: dict[str, str], body: bytes | None) -> dict[str, str]: """The headers actually sent upstream. A signing mount gets a signature minted over the upstream URL, because the edge rewrote the Host the proxy @@ -511,20 +567,21 @@ class CacheEdge: if isinstance(prepared, NetworkError): self.reject(mount, UNREACHABLE) return prepared - identity: Final = request_identity( - self.secret, test_key, method, url, self.keyed(mount, prepared.headers), body, - ) + keyed_headers: Final = self.keyed(mount, prepared.headers) + identity: Final = request_identity(self.secret, test_key, method, url, keyed_headers, body) key: Final = slotted_key(self.secret, identity, self.slots.take(identity)) found: Final = self.lookup(key) if isinstance(found, CacheHit): response: Final = decode_response(self.secret, key, found.payload, mount, url) if response is not None and self.clock() < found.valid_until: self.count(mount, "hits") + self.record_key(mount, "hit", test_key, method, url, keyed_headers, body) return StreamHead(response.status_code, response.headers, response_steps(response)) self.count(mount, "corrupt" if response is None else "expired") self.store.discard(key, found.payload) capture_slot: Final = self.lookup(key) if isinstance(found, CacheHit) else found self.count(mount, "misses") + self.record_key(mount, "miss", test_key, method, url, keyed_headers, body) if isinstance(capture_slot, CacheUnavailable): self.count(mount, "cache_errors") self.count(mount, "upstream_attempts") diff --git a/tests/e2e/provider_cache_redis.py b/tests/e2e/provider_cache_redis.py index be4e31b2c49..2c7419cfc0f 100644 --- a/tests/e2e/provider_cache_redis.py +++ b/tests/e2e/provider_cache_redis.py @@ -134,6 +134,10 @@ def write_metrics(cache: CacheEdge) -> None: root: Final = Path(directory) root.mkdir(parents=True, exist_ok=True) (root / f"{os.getpid()}.json").write_text(report + "\n") + if cache.probe.rows: + (root / f"keys-{os.getpid()}.json").write_text( + json.dumps([dict(row) for row in cache.probe.rows]) + "\n" + ) except OSError: logging.getLogger(__name__).warning("provider cache metrics artifact unavailable") logging.getLogger(__name__).info("%s", report) From ea5887d4bb5eec62dfab82b2f71183e210cf75b3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:10:48 -0700 Subject: [PATCH 136/187] test: cover proxy_admin_viewer sessions in the model_group/info admin regression test --- .../proxy/proxy_server/test_routes_model_info.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index a076b3593b1..2c101156c6c 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -340,10 +340,13 @@ def model_group_info_router(monkeypatch): return router -def test_model_group_info_proxy_admin_ignores_key_model_restriction(client, auth_as, model_group_info_router): +@pytest.mark.parametrize("admin_role", ["proxy_admin", "proxy_admin_viewer"]) +def test_model_group_info_proxy_admin_ignores_key_model_restriction( + client, auth_as, model_group_info_router, admin_role +): from litellm.proxy._types import LitellmUserRoles - with auth_as(LitellmUserRoles.PROXY_ADMIN, models=["no-default-models"]): + with auth_as(LitellmUserRoles(admin_role), models=["no-default-models"]): response = client.get("/model_group/info") assert response.status_code == 200 From a9fc6d255b5bd82220f3c2265eaf568dcf179423 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 15:15:42 -0700 Subject: [PATCH 137/187] fix(prices): dedupe Nova cache_read_input_token_cost keys left by a text merge Six Amazon Nova entries define cache_read_input_token_cost twice, which is what a clean text merge of two branches that both added the field looks like. JSON parsers keep the last occurrence, so this turned test_price_map_has_no_duplicate_keys red on every open PR's merge commit, including this one, which touches neither file. Both occurrences in all six entries carry the same value, so dropping the later one leaves every parsed price identical. Same change as #41496, carried here so this branch is not blocked on it. Identical deletions, so the two merge cleanly in either order. --- ...model_prices_and_context_window_backup.json | 18 ++++++------------ model_prices_and_context_window.json | 18 ++++++------------ 2 files changed, 12 insertions(+), 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b1c38d0350a..92e5b1c4ff7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -377,8 +377,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -561,8 +560,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -578,8 +576,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45794,8 +45791,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45809,8 +45805,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45842,8 +45837,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b1c38d0350a..92e5b1c4ff7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -377,8 +377,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -561,8 +560,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -578,8 +576,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45794,8 +45791,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 1.5e-08 + "supports_tool_choice": true }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45809,8 +45805,7 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 8.75e-09 + "supports_tool_choice": true }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45842,8 +45837,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true, - "cache_read_input_token_cost": 2e-07 + "supports_tool_choice": true }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, From b9751a38ab0a21f56703569b18b3db084a6fa1ae Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 15:36:54 -0700 Subject: [PATCH 138/187] Revert "fix(prices): dedupe Nova cache_read_input_token_cost keys left by a text merge" This reverts commit a9fc6d255b5bd82220f3c2265eaf568dcf179423. --- ...model_prices_and_context_window_backup.json | 18 ++++++++++++------ model_prices_and_context_window.json | 18 ++++++++++++------ 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 92e5b1c4ff7..b1c38d0350a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -377,7 +377,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -560,7 +561,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -576,7 +578,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45791,7 +45794,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45805,7 +45809,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45837,7 +45842,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 92e5b1c4ff7..b1c38d0350a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -377,7 +377,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "amazon.nova-2-lite-v1:0": { "cache_read_input_token_cost": 7.5e-08, @@ -560,7 +561,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -576,7 +578,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "amazon.nova-sonic-v1:0": { "deprecation_date": "2026-09-14", @@ -45791,7 +45794,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 1.5e-08 }, "us.amazon.nova-micro-v1:0": { "cache_read_input_token_cost": 8.75e-09, @@ -45805,7 +45809,8 @@ "supports_function_calling": true, "supports_prompt_caching": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 8.75e-09 }, "us.amazon.nova-premier-v1:0": { "deprecation_date": "2026-09-14", @@ -45837,7 +45842,8 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_vision": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "cache_read_input_token_cost": 2e-07 }, "us.anthropic.claude-3-5-haiku-20241022-v1:0": { "cache_creation_input_token_cost": 1e-06, From 9b77b5c2cb5ec8bbf58738d279e266e4f8162212 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 5 Sep 2026 18:08:09 -0700 Subject: [PATCH 139/187] feat(proxy): let proxy admins choose which team fields team admins may edit Team admins could never reach POST /team/update: the route gate answered 401 before the handler's team-admin branch ran. This moves /team/update into the self-managed routes, resolves proxy admin, org admin or team admin inside the handler, and filters team admins through a new proxy-wide UI setting, team_admin_editable_team_fields. The setting is an allow-list of team fields. Empty means team admins cannot edit team settings and get a 403 pointing at the proxy admin, and changing a field outside the list fails 403 naming that field. Only values that differ from what is stored count, since the dashboard resends the whole form. The registry of fields the setting accepts ships empty on purpose. Each field lands in its own follow-up PR with its value diff and dashboard wiring. The Admin UI gains a "Team admin editable fields" section under Settings > UI and a toast on the team page while editing is disabled. Refs LIT-5722 Claude-Session: https://claude.ai/code/session_01A6SkwJdfZUmkzfUkrEkqX8 --- litellm/proxy/_types.py | 3 + .../team_admin_field_permissions.py | 179 ++++++++++++++ .../management_endpoints/team_endpoints.py | 84 +++++-- .../proxy_setting_endpoints.py | 35 +++ .../management/test_team_update.py | 63 ++--- .../proxy/auth/test_route_checks.py | 83 ++++--- .../test_team_admin_field_permissions.py | 125 ++++++++++ .../test_team_endpoints.py | 234 +++++++++++++++++- .../test_proxy_setting_endpoints.py | 108 ++++++++ .../TeamAdminEditableFieldsSettings.test.tsx | 91 +++++++ .../TeamAdminEditableFieldsSettings.tsx | 64 +++++ .../UISettings/UISettings.test.tsx | 44 ++++ .../AdminSettings/UISettings/UISettings.tsx | 26 ++ .../src/components/team/TeamInfo.test.tsx | 65 +++++ .../src/components/team/TeamInfo.tsx | 27 +- .../team/teamAdminEditAccess.test.ts | 69 ++++++ .../components/team/teamAdminEditAccess.ts | 38 +++ 17 files changed, 1244 insertions(+), 94 deletions(-) create mode 100644 litellm/proxy/management_endpoints/team_admin_field_permissions.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx create mode 100644 ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts create mode 100644 ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 14e3635f079..94cfd090dcc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -844,6 +844,9 @@ class LiteLLMRoutes(enum.Enum): ) self_managed_routes = [ + # update_team resolves proxy/org/team admin itself and filters team admins + # through the team_admin_editable_team_fields setting + "/team/update", "/team/member_add", "/team/member_delete", "/management/v1/teams/{team_id}/members/bulk_delete", diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py new file mode 100644 index 00000000000..9ccacf9c54f --- /dev/null +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -0,0 +1,179 @@ +"""Proxy-wide allow-list of team-settings fields a team admin may change on /team/update.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal, TypeAlias, assert_never + +from fastapi import HTTPException +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.models.team import LiteLLM_TeamTable +from litellm.proxy._types import ( + LiteLLM_ManagementEndpoint_MetadataFields, + LiteLLM_ManagementEndpoint_MetadataFields_Premium, + UpdateTeamRequest, +) + +TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_fields" + +# TODO(LIT-5722): stays empty until each field's value-diff and dashboard wiring lands, one field per PR +SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset() + +_FIELD_LIST: Final = TypeAdapter(list[str]) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) +_METADATA_FOLDED_FIELDS: Final[frozenset[str]] = frozenset( + (*LiteLLM_ManagementEndpoint_MetadataFields, *LiteLLM_ManagementEndpoint_MetadataFields_Premium) +) +_SYSTEM_MANAGED_METADATA_KEYS: Final[frozenset[str]] = frozenset({"team_member_budget_id"}) +_NOT_COLUMNS: Final[frozenset[str]] = frozenset({"team_id", "metadata"}) +_SETTINGS_LOCATION: Final = "Settings > UI > Team admin editable fields" + + +@dataclass(frozen=True, slots=True) +class TeamAdminEditAllowed: + kind: Literal["allowed"] = "allowed" + + +@dataclass(frozen=True, slots=True) +class TeamAdminEditingDisabled: + kind: Literal["disabled"] = "disabled" + + +@dataclass(frozen=True, slots=True) +class TeamAdminFieldNotPermitted: + field: str + kind: Literal["field_not_permitted"] = "field_not_permitted" + + +TeamAdminEditVerdict: TypeAlias = TeamAdminEditAllowed | TeamAdminEditingDisabled | TeamAdminFieldNotPermitted + + +def resolve_team_admin_editable_fields( + general_settings: Mapping[str, object], + supported: frozenset[str], +) -> frozenset[str]: + raw: Final = general_settings.get(TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING) + if raw is None: + return frozenset() + try: + configured: Final = frozenset(_FIELD_LIST.validate_python(raw)) + except ValidationError: + verbose_proxy_logger.warning( + "%s must be a list of field names; ignoring %r", TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, raw + ) + return frozenset() + unsupported: Final = configured - supported + if unsupported: + verbose_proxy_logger.warning( + "%s ignores unsupported field(s) %s; supported: %s", + TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, + sorted(unsupported), + sorted(supported), + ) + return configured & supported + + +def _as_object(value: object) -> Mapping[str, object]: + try: + return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value) + except ValidationError: + return _EMPTY + + +def _stored_metadata(existing: Mapping[str, object]) -> Mapping[str, object]: + return _as_object(existing.get("metadata")) + + +def _submitted_metadata( + data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object] +) -> Mapping[str, object]: + """Metadata as it would be stored: the caller's dict (or the stored one) with top-level folded fields laid over.""" + base: Final = ( + _as_object(submitted.get("metadata")) if "metadata" in data.model_fields_set else _stored_metadata(existing) + ) + folded: Final = data.model_fields_set & _METADATA_FOLDED_FIELDS + return MappingProxyType({key: submitted[key] if key in folded else base[key] for key in base.keys() | folded}) + + +def _metadata_changes( + data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object] +) -> frozenset[str]: + merged: Final = _submitted_metadata(data, submitted, existing) + stored: Final = _stored_metadata(existing) + return frozenset( + key if key in _METADATA_FOLDED_FIELDS else "metadata" + for key in (merged.keys() | stored.keys()) - _SYSTEM_MANAGED_METADATA_KEYS + if merged.get(key) != stored.get(key) + ) + + +def _stored_model_aliases(existing_row: LiteLLM_TeamTable) -> Mapping[str, object]: + table: Final = existing_row.litellm_model_table + return _as_object(_JSON_OBJECT.validate_json(table.model_dump_json()).get("model_aliases")) if table else _EMPTY + + +def _column_changed( + field: str, submitted: Mapping[str, object], existing: Mapping[str, object], existing_row: LiteLLM_TeamTable +) -> bool: + if field == "model_aliases": + return _as_object(submitted.get(field)) != _stored_model_aliases(existing_row) + if field in LiteLLM_TeamTable.model_fields: + return submitted.get(field) != existing.get(field) + return True + + +def changed_team_fields(data: UpdateTeamRequest, existing_row: LiteLLM_TeamTable) -> frozenset[str]: + """Logical field names whose stored value the request would change. + + Request and stored row are compared as JSON values so both sides share one representation. Fields the + server folds into metadata are attributed to their own name whether they arrive top-level or inside + ``metadata``; anything else in ``metadata`` is attributed to ``metadata``. Fields with no stored + counterpart on the team row count as changed whenever they are sent. + """ + submitted: Final = _JSON_OBJECT.validate_json(data.model_dump_json(exclude_unset=True)) + existing: Final = _JSON_OBJECT.validate_json(existing_row.model_dump_json()) + column_fields: Final = frozenset(data.model_fields_set) - _NOT_COLUMNS - _METADATA_FOLDED_FIELDS + column_changes: Final = frozenset( + field for field in column_fields if _column_changed(field, submitted, existing, existing_row) + ) + return column_changes | _metadata_changes(data, submitted, existing) + + +def team_admin_edit_verdict( + data: UpdateTeamRequest, + existing: LiteLLM_TeamTable, + permitted: frozenset[str], +) -> TeamAdminEditVerdict: + if not permitted: + return TeamAdminEditingDisabled() + blocked: Final = sorted(changed_team_fields(data, existing) - permitted) + if blocked: + return TeamAdminFieldNotPermitted(field=blocked[0]) + return TeamAdminEditAllowed() + + +def raise_for_team_admin_edit_verdict(verdict: TeamAdminEditVerdict) -> None: + match verdict: + case TeamAdminEditAllowed(): + return + case TeamAdminEditingDisabled(): + raise HTTPException( + status_code=403, + detail=( + "Team admins on this proxy cannot edit team settings. " + f"Ask a proxy admin to enable fields under {_SETTINGS_LOCATION}." + ), + ) + case TeamAdminFieldNotPermitted(field=field): + raise HTTPException( + status_code=403, + detail=( + f"Team admins on this proxy do not have permission to update '{field}'. " + f"Ask a proxy admin to add it under {_SETTINGS_LOCATION}." + ), + ) + case _: + assert_never(verdict) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 23b02c22f24..d88d6d13720 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -18,7 +18,18 @@ from collections.abc import Iterable, Mapping, Sequence from collections.abc import Set as AbstractSet from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast +from typing import ( + TYPE_CHECKING, + Annotated, + Final, + Literal, + NamedTuple, + NoReturn, + Protocol, + TypeAlias, + TypeVar, + cast, +) import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -122,6 +133,12 @@ from litellm.proxy.management_endpoints.router_weights import validate_router_se from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, ) +from litellm.proxy.management_endpoints.team_admin_field_permissions import ( + SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS, + raise_for_team_admin_edit_verdict, + resolve_team_admin_editable_fields, + team_admin_edit_verdict, +) from litellm.proxy.management_helpers.access_group_team_sync import ( TEAM_ADVISORY_LOCK_SQL, AccessGroupSyncTx, @@ -439,32 +456,43 @@ async def _refresh_cached_team( ) -async def _can_manage_team( +TeamAccessRole: TypeAlias = Literal["proxy_admin", "org_admin", "team_admin"] + + +def _raise_team_access_denied() -> NoReturn: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="You do not have access to this team", + ) + + +async def _resolve_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, -) -> bool: - """True for a proxy admin, an admin of this team, or an org admin for the team's organization.""" +) -> TeamAccessRole | None: + """Strongest role the caller holds over ``team_obj``, or None when they hold none. + + Org admin outranks team admin so a caller holding both keeps unrestricted edits. + """ if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: - return True + return "proxy_admin" + + if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj): + return "org_admin" if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj): - return True + return "team_admin" - return await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj) + return None async def _verify_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, ) -> None: - """Raise HTTPException(403) unless the caller can manage the given team.""" - if await _can_manage_team(team_obj=team_obj, user_api_key_dict=user_api_key_dict): - return - - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail="You do not have access to this team", - ) + """Raise 403 unless the caller is a proxy admin, an org admin for the team's org, or a team admin.""" + if await _resolve_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) is None: + _raise_team_access_denied() class TeamMemberBudgetHandler: @@ -2096,6 +2124,7 @@ async def update_team( try: from litellm.proxy.management_helpers.audit_logs import is_audit_logging_enabled from litellm.proxy.proxy_server import ( + general_settings, litellm_proxy_admin_name, llm_router, premium_user, @@ -2144,16 +2173,29 @@ async def update_team( ) if existing_team_row is None: + # Non-proxy-admins get the same 403 as an access denial so /team/update + # cannot be used to probe which team ids exist + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + _raise_team_access_denied() raise HTTPException( status_code=404, detail={"error": f"Team not found, passed team_id={data.team_id}"}, ) - # Verify caller has access to manage this team - await _verify_team_access( - team_obj=LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()), - user_api_key_dict=user_api_key_dict, - ) + existing_team: Final = LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()) + access_role: Final = await _resolve_team_access(team_obj=existing_team, user_api_key_dict=user_api_key_dict) + if access_role is None: + _raise_team_access_denied() + if access_role == "team_admin": + raise_for_team_admin_edit_verdict( + team_admin_edit_verdict( + data=data, + existing=existing_team, + permitted=resolve_team_admin_editable_fields( + general_settings, SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS + ), + ) + ) await validate_router_settings_weights( data.router_settings, @@ -4585,7 +4627,7 @@ async def team_info( await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_table) organization_models: Final[list[str] | None] = ( _parent_organization_models(team_info) - if await _can_manage_team(team_obj=team_table, user_api_key_dict=user_api_key_dict) + if await _resolve_team_access(team_obj=team_table, user_api_key_dict=user_api_key_dict) is not None else None ) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index c12d071dd36..c69225c55d6 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -29,6 +29,10 @@ from litellm.proxy.config_resolvers.sso import ( SSO_SECRET_FIELDS, resolve_sso_config, ) +from litellm.proxy.management_endpoints.team_admin_field_permissions import ( + SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS, + TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, +) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import invalidate_config_param from litellm.repositories.config_repository import ConfigRepository @@ -212,6 +216,9 @@ class UIThemeSettingsResponse(SettingsResponse): """Response model for UI theme settings""" +_TEAM_ADMIN_FIELD_ENUM: Final = tuple(sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS)) + + class UISettings(BaseModel): """Configuration for UI-specific flags""" @@ -304,6 +311,18 @@ class UISettings(BaseModel): description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.", ) + team_admin_editable_team_fields: Sequence[str] = Field( + default=(), + description=( + "Team settings fields a team admin may change on the teams they administer. " + "Empty means team admins cannot edit team settings at all. " + "Proxy admins and org admins are not affected." + ), + json_schema_extra={ # mutable-ok: pydantic only merges json_schema_extra when it is a plain dict + "items": {"type": "string", "enum": [*_TEAM_ADMIN_FIELD_ENUM]}, # mutable-ok: nested in the dict above + }, + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" @@ -326,6 +345,7 @@ ALLOWED_UI_SETTINGS_FIELDS: Final = { "disable_custom_api_keys", "disable_key_generate_for_org_admin", "enable_chat_ui", + TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, } ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: Final = "enable_ptu_cost_attribution" @@ -360,6 +380,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [ "disable_vector_stores_for_internal_users", "allow_vector_stores_for_team_admins", "disable_key_generate_for_org_admin", + TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, ] # Extension point: packages outside OSS (e.g. litellm_enterprise) can @@ -1571,6 +1592,20 @@ async def update_ui_settings( except ValidationError as e: raise HTTPException(status_code=422, detail=e.errors()) + unsupported_team_fields: Final = sorted( + frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS + ) + if unsupported_team_fields: + raise HTTPException( + status_code=400, + detail={ # mutable-ok: HTTPException detail must be a plain dict for FastAPI JSON serialization + "error": ( + f"{TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING} does not support {unsupported_team_fields}. " + f"Supported fields: {sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS)}." + ) + }, + ) + # Only include fields the caller actually sent (not Pydantic defaults). settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True) diff --git a/tests/proxy_behavior/management/test_team_update.py b/tests/proxy_behavior/management/test_team_update.py index 23ea89fa74d..d6273e64132 100644 --- a/tests/proxy_behavior/management/test_team_update.py +++ b/tests/proxy_behavior/management/test_team_update.py @@ -9,31 +9,31 @@ pytestmark = pytest.mark.asyncio(loop_scope="session") # POST /team/update — actor x team-shape matrix (shapes built by _seed_target). -# Each request carries the team's own organization_id so a non-proxy-admin can -# reach the org-scoped branch of the route-permission gate (401 on denial), -# which fronts the handler's _verify_team_access. Only PROXY_ADMIN and an -# ORG_ADMIN of the team's org pass: an internal_user team admin is filtered by -# the route gate before _verify_team_access's team-admin branch is reached. +# The route is self-managed (LIT-5722), so every authenticated caller reaches +# update_team and denials are the handler's 403, never the route gate's 401. +# Only PROXY_ADMIN and an ORG_ADMIN of the team's org pass: a team admin is +# admitted by _resolve_team_access but then refused because no team field is +# enabled for team admins (team_admin_editable_team_fields ships empty). MARKER_ALIAS = "behavior-pin-update-marker-alias" _MATRIX = [ ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), - ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 401), - ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 401), - ("alpha/owner", Actor.OWNER, "alpha", 401), - ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 401), - ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 401), - ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 401), - ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 401), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 403), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), - ("beta/org_admin", Actor.ORG_ADMIN, "beta", 401), - ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 401), - ("beta/internal_user", Actor.INTERNAL_USER, "beta", 401), - ("beta/owner", Actor.OWNER, "beta", 401), - ("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 401), - ("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 401), - ("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 401), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/internal_user", Actor.INTERNAL_USER, "beta", 403), + ("beta/owner", Actor.OWNER, "beta", 403), + ("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 403), + ("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 403), + ("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 403), ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), ] @@ -110,8 +110,9 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context( ): """With no organization_id in the body the route gate resolves the target team's org from team_id, so an org admin of the team's own org is allowed - (200), same as PROXY_ADMIN. A team admin of that same team stays denied - (401): the resolution grants org admins access, not team admins.""" + (200), same as PROXY_ADMIN. A team admin of that same team reaches the + handler but is refused (403) until a proxy admin enables fields for team + admins, and the response says so.""" await _seed_target(prisma, world, "alpha", scratch.prefix) allowed_org_admin = await proxy_client.post( @@ -133,21 +134,25 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context( headers={"Authorization": f"Bearer {world.keys[Actor.TEAM_ADMIN].cleartext}"}, json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS}, ) - assert denied_team_admin.status_code == 401, denied_team_admin.text + assert denied_team_admin.status_code == 403, denied_team_admin.text + assert "cannot edit team settings" in denied_team_admin.text, denied_team_admin.text + assert "Team admin editable fields" in denied_team_admin.text, denied_team_admin.text # Relocation gate — moving a team to a different org. The scratch team starts # in ORG_A; each scenario relocates it to ORG_B. PROXY_ADMIN bypasses; -# ORG_B_ADMIN clears the route gate (dest-org admin) but fails -# _verify_team_access on the source team (403); the rest fail the route gate -# (401). The relocation-*allowed* branch (caller is org admin of both orgs) is -# covered by test_team_update_org_relocation_allowed_for_dual_org_admin below. +# ORG_B_ADMIN reaches the handler but holds no role on the source team (403); +# ORG_ADMIN holds the source team but not the destination org (403 from the +# relocation gate); the team admin is refused by the empty field allow-list and +# the internal user holds no role at all (403). The relocation-*allowed* branch +# (caller is org admin of both orgs) is covered by +# test_team_update_org_relocation_allowed_for_dual_org_admin below. _RELOCATION = [ ("proxy_admin", Actor.PROXY_ADMIN, 200), ("org_b_admin", Actor.ORG_B_ADMIN, 403), - ("org_admin", Actor.ORG_ADMIN, 401), - ("team_admin", Actor.TEAM_ADMIN, 401), - ("internal_user", Actor.INTERNAL_USER, 401), + ("org_admin", Actor.ORG_ADMIN, 403), + ("team_admin", Actor.TEAM_ADMIN, 403), + ("internal_user", Actor.INTERNAL_USER, 403), ] diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 806c55d51ce..ed319190a6b 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -2892,45 +2892,53 @@ def test_team_update_gate_allows_org_admin_with_resolved_org(): ) -def test_team_update_gate_rejects_without_org_context(): - """Without organization_id (i.e. resolution found no org, or a non-org-admin), - the gate still rejects /team/update — the fix adds no blanket allow. Guards - against re-widening the route (e.g. dropping it into self_managed_routes).""" +def test_team_update_gate_admits_internal_user_without_org_context(): + """/team/update is self-managed (LIT-5722): the coarse gate admits any authenticated + caller and update_team resolves proxy, org or team admin itself, then filters team admins + through the team_admin_editable_team_fields setting. Before that the gate 401'd every + team admin, which left the handler's team-admin branch unreachable.""" + from litellm.proxy._types import LiteLLMRoutes + + assert "/team/update" in LiteLLMRoutes.self_managed_routes.value + + user_obj = LiteLLM_UserTable( + user_id="team-admin-user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + organization_memberships=None, + ) + valid_token = UserAPIKeyAuth(user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value) + request = MagicMock(spec=Request) + request.method = "POST" + request.query_params = {} + + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/team/update", + request=request, + valid_token=valid_token, + request_data={"team_id": "team-1", "max_budget": 42}, + ) + + +def test_team_update_gate_defers_cross_org_admin_to_the_handler(): # test-quality-ok: the gate's only success signal is not raising; the handler's 403 it defers to is pinned in test_team_endpoints + """An org admin of a DIFFERENT org clears the coarse gate like any internal user; + update_team's _resolve_team_access finds no role on the team and 403s (pinned in + test_team_endpoints), so there is still no cross-org escalation.""" user_obj = _make_org_admin_user("org-1") valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value) request = MagicMock(spec=Request) request.method = "POST" request.query_params = {} - with pytest.raises(Exception, match="Only proxy admin can be used to generate"): - RouteChecks.non_proxy_admin_allowed_routes_check( - user_obj=user_obj, - _user_role=LitellmUserRoles.INTERNAL_USER.value, - route="/team/update", - request=request, - valid_token=valid_token, - request_data={"team_id": "team-1", "max_budget": 42}, - ) - - -def test_team_update_gate_rejects_cross_org_admin_with_resolved_org(): - """Even after the target team's org is resolved, an org admin of a DIFFERENT - org is rejected at the gate (no cross-org escalation).""" - user_obj = _make_org_admin_user("org-1") - valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value) - request = MagicMock(spec=Request) - request.method = "POST" - request.query_params = {} - - with pytest.raises(Exception, match="Only proxy admin can be used to generate"): - RouteChecks.non_proxy_admin_allowed_routes_check( - user_obj=user_obj, - _user_role=LitellmUserRoles.INTERNAL_USER.value, - route="/team/update", - request=request, - valid_token=valid_token, - request_data={"team_id": "team-1", "organization_id": "org-2"}, - ) + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/team/update", + request=request, + valid_token=valid_token, + request_data={"team_id": "team-1", "organization_id": "org-2"}, + ) # ── PATCH /team/{team_id}: same org-context + role reach as POST /team/update ── @@ -2993,10 +3001,11 @@ async def test_add_team_org_context_noop_for_static_team_route(): assert out == body -def test_patch_team_route_has_same_reach_as_team_update(): - """/team/{team_id} is reachable by org admins (in org_admin_allowed_routes) but - NOT by regular internal users or the role-agnostic self_managed_routes — the - latter would open /team/new (the collision footgun) to any authenticated user.""" +def test_patch_team_route_stays_out_of_self_managed_routes(): + """Unlike POST /team/update, PATCH /team/{team_id} cannot be self-managed: its + template also matches /team/new (the collision footgun), so it stays reachable by + org admins (org_admin_allowed_routes) and proxy admins only, never by regular + internal users or through the role-agnostic self_managed_routes.""" from litellm.proxy._types import LiteLLMRoutes assert RouteChecks.check_route_access( diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py new file mode 100644 index 00000000000..91479921c61 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -0,0 +1,125 @@ +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LiteLLM_ModelTable, LiteLLM_TeamTable, UpdateTeamRequest +from litellm.proxy.management_endpoints.team_admin_field_permissions import ( + TeamAdminEditAllowed, + TeamAdminEditingDisabled, + TeamAdminFieldNotPermitted, + changed_team_fields, + raise_for_team_admin_edit_verdict, + resolve_team_admin_editable_fields, + team_admin_edit_verdict, +) + +_SUPPORTED = frozenset({"tpm_limit", "rpm_limit", "team_alias"}) + + +def _team(**overrides): + return LiteLLM_TeamTable(team_id="team-1", **overrides) + + +class TestResolveTeamAdminEditableFields: + def test_missing_setting_means_nothing_editable(self): + assert resolve_team_admin_editable_fields({}, _SUPPORTED) == frozenset() + + def test_keeps_only_supported_names(self): + configured = {"team_admin_editable_team_fields": ["tpm_limit", "blocked", "organization_id"]} + assert resolve_team_admin_editable_fields(configured, _SUPPORTED) == frozenset({"tpm_limit"}) + + @pytest.mark.parametrize("raw", ["tpm_limit", 7, {"tpm_limit": True}, [1, 2]]) + def test_malformed_setting_fails_closed(self, raw): + assert resolve_team_admin_editable_fields({"team_admin_editable_team_fields": raw}, _SUPPORTED) == frozenset() + + +class TestChangedTeamFields: + def test_team_id_alone_changes_nothing(self): + assert changed_team_fields(UpdateTeamRequest(team_id="team-1"), _team()) == frozenset() + + def test_column_echoing_stored_value_is_not_a_change(self): + data = UpdateTeamRequest(team_id="team-1", tpm_limit=5, team_alias="alpha", max_budget=None) + assert changed_team_fields(data, _team(tpm_limit=5, team_alias="alpha")) == frozenset() + + def test_column_with_different_value_is_a_change(self): + data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha") + assert changed_team_fields(data, _team(tpm_limit=5, team_alias="alpha")) == frozenset({"tpm_limit"}) + + def test_explicit_null_clearing_a_stored_column_is_a_change(self): + data = UpdateTeamRequest(team_id="team-1", max_budget=None) + assert changed_team_fields(data, _team(max_budget=30.0)) == frozenset({"max_budget"}) + + def test_folded_field_sent_top_level_is_named_not_metadata(self): + data = UpdateTeamRequest(team_id="team-1", guardrails=["b"]) + assert changed_team_fields(data, _team(metadata={"guardrails": ["a"]})) == frozenset({"guardrails"}) + + def test_folded_field_sent_inside_metadata_is_named_not_metadata(self): + data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["b"]}) + assert changed_team_fields(data, _team(metadata={"guardrails": ["a"]})) == frozenset({"guardrails"}) + + def test_custom_metadata_key_change_is_attributed_to_metadata(self): + data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["a"], "cost_center": "b"}) + existing = _team(metadata={"guardrails": ["a"], "cost_center": "a"}) + assert changed_team_fields(data, existing) == frozenset({"metadata"}) + + def test_metadata_echo_with_top_level_override_only_names_the_override(self): + data = UpdateTeamRequest(team_id="team-1", guardrails=["b"], metadata={"guardrails": ["a"], "cost_center": "a"}) + existing = _team(metadata={"guardrails": ["a"], "cost_center": "a"}) + assert changed_team_fields(data, existing) == frozenset({"guardrails"}) + + def test_dropping_a_stored_key_from_submitted_metadata_is_a_change(self): + data = UpdateTeamRequest(team_id="team-1", metadata={"cost_center": "a"}) + existing = _team(metadata={"cost_center": "a", "tags": ["x"], "logging": [{"callback": "langfuse"}]}) + assert changed_team_fields(data, existing) == frozenset({"tags", "logging"}) + + def test_server_managed_metadata_key_is_ignored(self): + data = UpdateTeamRequest(team_id="team-1", metadata={"cost_center": "a"}) + existing = _team(metadata={"cost_center": "a", "team_member_budget_id": "budget-1"}) + assert changed_team_fields(data, existing) == frozenset() + + def test_model_aliases_compare_against_the_model_table(self): + table = LiteLLM_ModelTable(model_aliases='{"fast": "gpt-4o-mini"}', created_by="a", updated_by="a") + same = UpdateTeamRequest(team_id="team-1", model_aliases={"fast": "gpt-4o-mini"}) + different = UpdateTeamRequest(team_id="team-1", model_aliases={"fast": "gpt-4o"}) + assert changed_team_fields(same, _team(litellm_model_table=table)) == frozenset() + assert changed_team_fields(different, _team(litellm_model_table=table)) == frozenset({"model_aliases"}) + + def test_empty_model_aliases_against_no_model_table_is_not_a_change(self): + assert changed_team_fields(UpdateTeamRequest(team_id="team-1", model_aliases={}), _team()) == frozenset() + + def test_field_without_a_stored_counterpart_counts_as_changed_when_sent(self): + data = UpdateTeamRequest(team_id="team-1", team_member_budget=10.0) + assert changed_team_fields(data, _team()) == frozenset({"team_member_budget"}) + + +class TestTeamAdminEditVerdict: + def test_no_permitted_fields_disables_editing_even_for_a_no_op(self): + verdict = team_admin_edit_verdict(UpdateTeamRequest(team_id="team-1"), _team(), frozenset()) + assert verdict == TeamAdminEditingDisabled() + + def test_changes_within_permitted_fields_are_allowed(self): + data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha") + verdict = team_admin_edit_verdict(data, _team(team_alias="alpha"), frozenset({"tpm_limit"})) + assert verdict == TeamAdminEditAllowed() + + def test_first_blocked_field_in_sorted_order_is_reported(self): + data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, rpm_limit=6, blocked=True) + verdict = team_admin_edit_verdict(data, _team(), frozenset({"tpm_limit"})) + assert verdict == TeamAdminFieldNotPermitted(field="blocked") + + +class TestRaiseForTeamAdminEditVerdict: + def test_allowed_does_not_raise(self): + assert raise_for_team_admin_edit_verdict(TeamAdminEditAllowed()) is None + + def test_disabled_is_a_403_pointing_at_the_proxy_admin(self): + with pytest.raises(HTTPException) as exc: + raise_for_team_admin_edit_verdict(TeamAdminEditingDisabled()) + assert exc.value.status_code == 403 + assert "cannot edit team settings" in exc.value.detail + assert "Settings > UI > Team admin editable fields" in exc.value.detail + + def test_field_not_permitted_is_a_403_naming_the_field(self): + with pytest.raises(HTTPException) as exc: + raise_for_team_admin_edit_verdict(TeamAdminFieldNotPermitted(field="blocked")) + assert exc.value.status_code == 403 + assert "'blocked'" in exc.value.detail diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 6478b18e553..733f6ddc36f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1,6 +1,6 @@ import asyncio import json -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, contextmanager from datetime import datetime, timezone from types import SimpleNamespace from typing import Final, Optional, cast @@ -76,6 +76,31 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( client = TestClient(app) +@contextmanager +def _team_admin_may_edit(*fields: str): + """Let team admins change ``fields`` on /team/update for the duration of the block. + + The registry ships empty (LIT-5722 adds fields one PR at a time), so tests that exercise the + gates layered underneath the allow-list widen it here instead of asserting the early 403.""" + with ( + patch( # test-quality-ok: the registry is a module constant update_team reads directly; no seam to inject + "litellm.proxy.management_endpoints.team_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS", + frozenset(fields), + ), + patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": list(fields)}), # test-quality-ok: update_team reads general_settings as a proxy_server module global + ): + yield + + +def _not_org_admin(): + """update_team asks whether the caller administers the team's org before it settles for team admin; + a MagicMock prisma cannot answer that lookup, so pin it to False.""" + return patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide + "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", + AsyncMock(return_value=False), + ) + + def _wire_team_create_tx(prisma_client): """`/team/new` inserts the team and mirrors it onto the access groups in one transaction, so a mocked client has to hand its team table back out of `db.tx()`. @@ -6393,6 +6418,7 @@ async def test_update_team_standalone_budget_raise_blocked_for_team_admin(): dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("max_budget"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6549,6 +6575,7 @@ async def test_update_team_standalone_budget_removal_blocked_for_team_admin(): dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("max_budget"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6618,6 +6645,7 @@ async def test_update_team_standalone_uncapped_team_admin_sets_finite_allowed( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("max_budget"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6712,6 +6740,7 @@ async def test_update_team_standalone_unchanged_budget_allowed( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("max_budget", "tpm_limit"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6810,6 +6839,7 @@ async def test_update_team_standalone_lower_budget_allowed( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("max_budget"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6912,6 +6942,8 @@ async def test_update_team_org_scoped_budget_exceeds_org_limit(): mock_org.litellm_budget_table = mock_budget_table with ( + _team_admin_may_edit("max_budget"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -6992,6 +7024,7 @@ async def test_update_team_standalone_models_not_gated_by_user_limit( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("models"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7091,6 +7124,8 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit( mock_org.litellm_budget_table = mock_budget_table with ( + _team_admin_may_edit("max_budget"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7202,6 +7237,8 @@ async def test_update_team_org_scoped_models_bypasses_user_limit( mock_org.litellm_budget_table = None with ( + _team_admin_may_edit("models"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7304,6 +7341,8 @@ async def test_update_team_org_scoped_models_not_in_org_models(): mock_org.litellm_budget_table = None with ( + _team_admin_may_edit("models"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7393,6 +7432,8 @@ async def test_update_team_org_scoped_models_with_all_proxy_models( mock_org.litellm_budget_table = None with ( + _team_admin_may_edit("models"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7502,6 +7543,7 @@ async def test_update_team_tpm_limit_not_gated_by_user_limit( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("tpm_limit"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7584,6 +7626,7 @@ async def test_update_team_rpm_limit_not_gated_by_user_limit( dummy_request = MagicMock(spec=Request) with ( + _team_admin_may_edit("rpm_limit"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -7981,6 +8024,8 @@ async def test_update_team_org_scoped_tpm_exceeds_org_limit(): mock_org.litellm_budget_table = mock_budget_table with ( + _team_admin_may_edit("tpm_limit"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -8067,6 +8112,8 @@ async def test_update_team_org_scoped_rpm_exceeds_org_limit(): mock_org.litellm_budget_table = mock_budget_table with ( + _team_admin_may_edit("rpm_limit"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -8158,6 +8205,8 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit( mock_org.litellm_budget_table = mock_budget_table with ( + _team_admin_may_edit("tpm_limit", "rpm_limit"), + _not_org_admin(), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -8286,6 +8335,7 @@ async def test_update_team_guardrails_with_org_id( } with ( + _team_admin_may_edit("guardrails", "organization_id"), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), @@ -11177,8 +11227,8 @@ async def test_update_team_blocks_non_admin_passthrough_routes(mock_db_client): mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing) with patch( - "litellm.proxy.management_endpoints.team_endpoints._verify_team_access", - AsyncMock(return_value=None), + "litellm.proxy.management_endpoints.team_endpoints._resolve_team_access", + AsyncMock(return_value="org_admin"), ): with pytest.raises(ProxyException) as exc: await update_team( @@ -13246,6 +13296,7 @@ async def test_update_team_output_token_estimate_lowered_rejected_for_team_admin with contextlib.ExitStack() as stack: _wire_update_team(stack, {_TEAM_ESTIMATE: 4000}) + stack.enter_context(_team_admin_may_edit("default_estimated_output_tokens")) with pytest.raises(ProxyException) as exc: await update_team( data=UpdateTeamRequest(team_id="test_team_id", default_estimated_output_tokens=1), @@ -13277,6 +13328,7 @@ async def test_update_team_output_token_estimate_unchanged_allows_team_admin_edi with contextlib.ExitStack() as stack: prisma = _wire_update_team(stack, {_TEAM_ESTIMATE: 4000}) + stack.enter_context(_team_admin_may_edit("team_alias")) await update_team( data=UpdateTeamRequest( team_id="test_team_id", @@ -13336,6 +13388,7 @@ async def test_update_team_batch_enqueued_token_limit_raised_rejected_for_team_a with contextlib.ExitStack() as stack: _wire_update_team(stack, {_TEAM_BATCH_LIMIT: 100000}) + stack.enter_context(_team_admin_may_edit("metadata")) with pytest.raises(ProxyException) as exc: await update_team( data=UpdateTeamRequest(team_id="test_team_id", metadata={_TEAM_BATCH_LIMIT: 10**12}), @@ -14894,3 +14947,178 @@ async def test_update_team_model_max_budget_raise_blocked_for_team_admin(): assert exc.value.code == "403" assert "proxy admin" in str(exc.value.message).lower() mock_prisma.db.litellm_teamtable.update.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# LIT-5722: team admins reach update_team through self_managed_routes and are +# filtered by the team_admin_editable_team_fields setting. +# --------------------------------------------------------------------------- + +_TEAM_ADMIN_CALLER = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-team-admin", user_id="team-admin" +) +_PROXY_ADMIN_CALLER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin") + + +def _update_request_stub(): + from unittest.mock import Mock + + from fastapi import Request + + return Mock(spec=Request) + + +@pytest.mark.asyncio +async def test_update_team_team_admin_is_refused_before_any_write_when_no_fields_are_enabled(): + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + stack.enter_context(_team_admin_may_edit()) + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert str(exc.value.code) == "403" + assert "cannot edit team settings" in str(exc.value.message) + assert not prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_update_team_configured_but_unsupported_field_does_not_open_editing(): + """Only fields in SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS count, whatever general_settings says.""" + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + stack.enter_context( + patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": ["team_alias"]}) # test-quality-ok: update_team reads general_settings as a proxy_server module global + ) + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert str(exc.value.code) == "403" + assert "cannot edit team settings" in str(exc.value.message) + assert not prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_update_team_team_admin_changing_an_unpermitted_field_is_refused_by_name(): + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + stack.enter_context(_team_admin_may_edit("team_alias")) + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed", tpm_limit=10), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert str(exc.value.code) == "403" + assert "'tpm_limit'" in str(exc.value.message) + assert not prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_update_team_team_admin_echoing_unpermitted_fields_unchanged_is_allowed( + disable_audit_logging_for_mocked_team, +): + """The dashboard resends the whole form, so only a value that differs from what is stored counts.""" + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + stack.enter_context(_team_admin_may_edit("team_alias")) + result = await update_team( + data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed", tpm_limit=None, models=[]), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert result["data"].team_id == "test_team_id" + assert prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_update_team_org_admin_is_not_filtered_by_the_team_admin_field_list( + disable_audit_logging_for_mocked_team, +): + """A caller who is both org admin and roster admin keeps unrestricted edits.""" + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + stack.enter_context(_team_admin_may_edit()) + stack.enter_context( + patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide + "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", + AsyncMock(return_value=True), + ) + ) + result = await update_team( + data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + assert result["data"].team_id == "test_team_id" + assert prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_update_team_unknown_team_is_403_for_non_proxy_admins_and_404_for_proxy_admins(): + """Now that any authenticated caller reaches the handler, 'team not found' must not leak team ids.""" + import contextlib + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + + with pytest.raises(ProxyException) as denied: + await update_team( + data=UpdateTeamRequest(team_id="no-such-team", team_alias="renamed"), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + with pytest.raises(ProxyException) as missing: + await update_team( + data=UpdateTeamRequest(team_id="no-such-team", team_alias="renamed"), + http_request=_update_request_stub(), + user_api_key_dict=_PROXY_ADMIN_CALLER, + ) + + assert str(denied.value.code) == "403" + assert "do not have access to this team" in str(denied.value.message) + assert "no-such-team" not in str(denied.value.message) + assert str(missing.value.code) == "404" + + +@pytest.mark.asyncio +async def test_resolve_team_access_ranks_proxy_admin_then_org_admin_then_team_admin(): + from litellm.proxy.management_endpoints.team_endpoints import _resolve_team_access + + team = LiteLLM_TeamTable( + team_id="team-1", + organization_id="org-1", + members_with_roles=[Member(user_id="team-admin", role="admin")], + ) + roster_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin") + outsider = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="someone-else") + org_lookup = AsyncMock(return_value=False) + + with patch("litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", org_lookup): # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide + assert await _resolve_team_access(team_obj=team, user_api_key_dict=_PROXY_ADMIN_CALLER) == "proxy_admin" + assert org_lookup.await_count == 0 + assert await _resolve_team_access(team_obj=team, user_api_key_dict=roster_admin) == "team_admin" + assert await _resolve_team_access(team_obj=team, user_api_key_dict=outsider) is None + org_lookup.return_value = True + assert await _resolve_team_access(team_obj=team, user_api_key_dict=roster_admin) == "org_admin" diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 709447d23c0..e846b8a5111 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3266,3 +3266,111 @@ class TestPtuCostAttributionUISetting: assert response.status_code == 400 assert "enable_ptu_cost_attribution" in str(response.json()["detail"]) assert not mock_prisma.db.litellm_uisettings.upsert.called + + +class TestTeamAdminEditableTeamFieldsSetting: + """team_admin_editable_team_fields: the proxy-wide allow-list update_team applies to team admins.""" + + def _as_proxy_admin(self, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + mock_prisma = MagicMock() + mock_prisma.db.litellm_uisettings.upsert = AsyncMock() + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + return mock_prisma + + def test_patch_rejects_field_names_the_proxy_does_not_support(self, monkeypatch): + mock_prisma = self._as_proxy_admin(monkeypatch) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS", + frozenset({"tpm_limit"}), + ) + + try: + response = client.patch( + "/update/ui_settings", + json={"team_admin_editable_team_fields": ["tpm_limit", "blocked", "organization_id"]}, + ) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 400 + detail = response.json()["detail"]["error"] + assert "['blocked', 'organization_id']" in detail + assert "['tpm_limit']" in detail + assert not mock_prisma.db.litellm_uisettings.upsert.called + + def test_patch_rejects_a_non_list_value(self, monkeypatch): + self._as_proxy_admin(monkeypatch) + + try: + response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": "tpm_limit"}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 422 + + def test_patch_persists_and_syncs_the_list_to_general_settings(self, monkeypatch): + mock_prisma = self._as_proxy_admin(monkeypatch) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS", + frozenset({"tpm_limit", "rpm_limit"}), + ) + general_settings: dict = {"team_admin_editable_team_fields": ["rpm_limit"]} + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + + try: + response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": ["tpm_limit"]}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"]) + assert stored["team_admin_editable_team_fields"] == ["tpm_limit"] + assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"] + + def test_patch_with_an_empty_list_turns_team_admin_editing_off_again(self, monkeypatch): + mock_prisma = self._as_proxy_admin(monkeypatch) + general_settings: dict = {"team_admin_editable_team_fields": ["tpm_limit"]} + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + + try: + response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": []}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"]) + assert stored["team_admin_editable_team_fields"] == [] + assert general_settings["team_admin_editable_team_fields"] == [] + + def test_get_reports_the_stored_list_and_advertises_supported_fields(self, mock_auth, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + mock_prisma = MagicMock() + mock_db_record = MagicMock() + mock_db_record.ui_settings = {"team_admin_editable_team_fields": ["tpm_limit"]} + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + general_settings: dict = {} + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + data = response.json() + assert data["values"]["team_admin_editable_team_fields"] == ["tpm_limit"] + assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"] + field_schema = data["field_schema"]["properties"]["team_admin_editable_team_fields"] + assert field_schema["type"] == "array" + assert field_schema["items"]["type"] == "string" + assert isinstance(field_schema["items"]["enum"], list) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx new file mode 100644 index 00000000000..d1a44460190 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx @@ -0,0 +1,91 @@ +import { describe, expect, it, vi } from "vitest"; +import userEvent from "@testing-library/user-event"; + +import { renderWithProviders, screen } from "@/../tests/test-utils"; + +import TeamAdminEditableFieldsSettings from "./TeamAdminEditableFieldsSettings"; + +describe("TeamAdminEditableFieldsSettings", () => { + it("explains that nothing can be enabled when the proxy supports no fields", () => { + renderWithProviders( + , + ); + + expect(screen.getByText("Team admins cannot edit team settings")).toBeInTheDocument(); + expect(screen.getByText(/does not support enabling any team settings fields/)).toBeInTheDocument(); + expect(screen.queryByRole("checkbox")).not.toBeInTheDocument(); + }); + + it("renders one checkbox per supported field, checked for the enabled ones", () => { + renderWithProviders( + , + ); + + expect(screen.getByText("1 field enabled")).toBeInTheDocument(); + expect(screen.getByText("Fields a team admin may change")).toBeInTheDocument(); + expect(screen.getByRole("checkbox", { name: "max_budget" })).not.toBeChecked(); + expect(screen.getByRole("checkbox", { name: "tpm_limit" })).toBeChecked(); + }); + + it("saves the list with the field added when an unchecked field is ticked", async () => { + const onUpdate = vi.fn(); + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(screen.getByRole("checkbox", { name: "max_budget" })); + + expect(onUpdate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["tpm_limit", "max_budget"] }); + }); + + it("saves the list with the field removed when a checked field is unticked", async () => { + const onUpdate = vi.fn(); + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(screen.getByRole("checkbox", { name: "tpm_limit" })); + + expect(onUpdate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["max_budget"] }); + }); + + it("blocks toggling while a save is in flight", async () => { + const onUpdate = vi.fn(); + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(screen.getByRole("checkbox", { name: "tpm_limit" })); + + expect(onUpdate).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx new file mode 100644 index 00000000000..8a42835ac4e --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx @@ -0,0 +1,64 @@ +"use client"; + +import { Badge } from "@/components/ui/badge"; +import { Checkbox } from "@/components/ui/checkbox"; + +interface TeamAdminEditableFieldsSettingsProps { + editableFields: readonly string[]; + supportedFields: readonly string[]; + description?: string; + isUpdating: boolean; + onUpdate: (settings: { team_admin_editable_team_fields: string[] }) => void; +} + +export default function TeamAdminEditableFieldsSettings({ + editableFields, + supportedFields, + description, + isUpdating, + onUpdate, +}: TeamAdminEditableFieldsSettingsProps) { + const toggleField = (field: string, checked: boolean) => { + const next = checked ? [...editableFields, field] : editableFields.filter((item) => item !== field); + onUpdate({ team_admin_editable_team_fields: next }); + }; + + return ( +

+
+
+

Team admin editable fields

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

{description}

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

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

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

Team admin editable fields

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

{description}

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

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

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

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

+ ); + } + + return ( + void submit(event)} className="space-y-4"> + ( +
+ {supportedFields.map((name) => { + const checkboxId = `team-admin-editable-${name}`; + return ( + + ); + })} +
+ )} + /> + + ); } diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx index f38caf0df81..c2834e65498 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx @@ -155,48 +155,4 @@ describe("UISettings", () => { ); expect(toast.success).toHaveBeenCalledWith("UI settings updated successfully"); }); - - it("saves the team admin editable field list when a supported field is ticked", () => { - const mutateMock = vi.fn((_settings, options) => { - options?.onSuccess?.(); - }); - mockUseUpdateUISettings.mockReturnValue({ - mutate: mutateMock, - isPending: false, - error: null, - }); - mockUseUISettings.mockReturnValue( - buildSettingsResponse({ - data: { - field_schema: { - properties: { - team_admin_editable_team_fields: { - description: "Team settings fields a team admin may change", - type: "array", - items: { type: "string", enum: ["tpm_limit"] }, - }, - }, - }, - values: { team_admin_editable_team_fields: [] }, - }, - }), - ); - - render(); - - expect(screen.getByText("Team settings fields a team admin may change")).toBeInTheDocument(); - - act(() => { - fireEvent.click(screen.getByRole("checkbox", { name: "Tokens per minute Limit (TPM)" })); - }); - - expect(mutateMock).toHaveBeenCalledWith( - { team_admin_editable_team_fields: ["tpm_limit"] }, - expect.objectContaining({ - onSuccess: expect.any(Function), - onError: expect.any(Function), - }), - ); - expect(toast.success).toHaveBeenCalledWith("Team admin editable fields updated successfully"); - }); }); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index 04c53ec39e8..612ca05d083 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -9,12 +9,7 @@ import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Separator } from "@/components/ui/separator"; import { Skeleton } from "@/components/ui/skeleton"; import { Switch } from "@/components/ui/switch"; -import { - parseSupportedTeamAdminEditableFields, - parseTeamAdminEditableFields, -} from "@/components/team/teamAdminEditAccess"; import PageVisibilitySettings from "./PageVisibilitySettings"; -import TeamAdminEditableFieldsSettings from "./TeamAdminEditableFieldsSettings"; interface SettingRowProps { ariaLabel: string; @@ -70,7 +65,6 @@ export default function UISettings() { const allowVectorStoresTeamAdminsProperty = schema?.properties?.allow_vector_stores_for_team_admins; const scopeUserSearchProperty = schema?.properties?.scope_user_search_to_org; const disableCustomApiKeysProperty = schema?.properties?.disable_custom_api_keys; - const teamAdminEditableFieldsProperty = schema?.properties?.team_admin_editable_team_fields; const values = data?.values ?? {}; const isDisabledForInternalUsers = Boolean(values.disable_model_add_for_internal_users); const isDisabledTeamAdminDeleteTeamUser = Boolean(values.disable_team_admin_delete_team_user); @@ -116,17 +110,6 @@ export default function UISettings() { }); }; - const handleUpdateTeamAdminEditableFields = (settings: { team_admin_editable_team_fields: string[] }) => { - updateSettings(settings, { - onSuccess: () => { - toast.success("Team admin editable fields updated successfully"); - }, - onError: (error) => { - toast.fromError(error); - }, - }); - }; - const handleToggleForwardClientHeaders = (checked: boolean) => { updateSettings( { forward_client_headers_to_llm_api: checked }, @@ -456,15 +439,6 @@ export default function UISettings() { isUpdating={isUpdating} onUpdate={handleUpdatePageVisibility} /> - - - )} From d9de2fc5f7cf23f78fed523f7704420a0d0cf680 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 14:38:41 -0700 Subject: [PATCH 148/187] style(ui): right-align the team admin editable fields Save button --- .../UISettings/TeamAdminEditableFieldsSettings.tsx | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx index 3737d30945e..e671a4a47b5 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.tsx @@ -128,9 +128,11 @@ function TeamAdminEditableFieldsForm({ )} /> - +
+ +
); } From cac521246a900637e43222e8d9db13886fd50eaf Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 14:58:22 -0700 Subject: [PATCH 149/187] style(ui): right-align the Save banner button --- .../AdminSettings/UserBannerSettings/UserBannerSettings.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings.tsx index 82c25150ec5..21e9ab0f7d9 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings.tsx @@ -145,7 +145,7 @@ function UserBannerSettingsForm({ persisted, isLoading, isPending, saveBanner }: )} -
+
From 287bbaa6c170191774f56d0143fb35fcda371a8b Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:08:08 +0000 Subject: [PATCH 150/187] fix(proxy): remove duplicate user budget hook that 429'd zero-cost models _PROXY_MaxBudgetLimiter re-checked spend:user:{id} against user_max_budget in async_pre_call_hook without the zero-cost model exemption that _user_max_budget_check applies in auth, so free models were rejected with "Max budget limit reached." once a user was over budget. Auth already owns this check, so the hook is deleted rather than taught the exemption again Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ARCHITECTURE.md | 3 +- .../common_utils/proxy_rate_limit_error.py | 2 +- litellm/proxy/hooks/__init__.py | 4 +- litellm/proxy/hooks/max_budget_limiter.py | 84 ------- litellm/proxy/utils.py | 4 +- .../budgets/BUDGET_TEST_COVERAGE_MATRIX.md | 2 +- .../test_unit_test_litellm_logging.py | 12 +- .../proxy/hooks/test_max_budget_limiter.py | 237 ------------------ .../test_proxy_rate_limit_provider_field.py | 61 +---- .../proxy/proxy_server/test_routes_config.py | 4 +- .../test_proxy_logging_hook_detection.py | 2 +- .../utils/proxy_logging/test_lifecycle.py | 10 +- .../utils/proxy_logging/test_pre_call_hook.py | 34 ++- .../test_rate_limit_error_unification.py | 69 ----- 14 files changed, 52 insertions(+), 476 deletions(-) delete mode 100644 litellm/proxy/hooks/max_budget_limiter.py delete mode 100644 tests/test_litellm/proxy/hooks/test_max_budget_limiter.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index b04e004aa1a..f418752d990 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -45,7 +45,7 @@ sequenceDiagram ProxyServer->>Auth: user_api_key_auth() Auth->>Redis: Check API key cache Redis-->>Auth: Key info + spend limits - ProxyServer->>Hooks: max_budget_limiter, parallel_request_limiter + ProxyServer->>Hooks: parallel_request_limiter, cache_control_check Hooks->>Redis: Check/increment rate limit counters ProxyServer->>Router: route_request() Router->>Main: litellm.acompletion() @@ -145,7 +145,6 @@ graph TD | Hook | File | Purpose | |------|------|---------| -| `max_budget_limiter` | `proxy/hooks/max_budget_limiter.py` | Enforce budget limits | | `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user | | `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation | | `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation | diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index c109da6f571..888a6d077ad 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -11,7 +11,7 @@ exception types: an upstream LLM provider returns 429. * :class:`fastapi.HTTPException` (status 429) — raised directly by proxy hooks such as ``parallel_request_limiter``, ``dynamic_rate_limiter``, - ``batch_rate_limiter``, ``max_budget_limiter``, ``max_iterations_limiter``, + ``batch_rate_limiter``, ``max_iterations_limiter``, etc. * :class:`litellm.llms.base_llm.chat.transformation.BaseLLMException` (status 429) — raised by some provider transports. diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index f3542098f95..a504c2ba102 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -4,7 +4,6 @@ from typing import Final, Literal from . import * from .cache_control_check import _PROXY_CacheControlCheck from .litellm_skills import SkillsInjectionHook -from .max_budget_limiter import _PROXY_MaxBudgetLimiter from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler from .max_iterations_limiter import _PROXY_MaxIterationsHandler from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler @@ -18,7 +17,6 @@ from .sensitive_data_routing import _PROXY_SensitiveDataRoutingHandler # transitively through `enterprise.enterprise_hooks` can resolve `PROXY_HOOKS` # and `get_proxy_hook` from this partially-initialized module without circling. PROXY_HOOKS: Final = { - "max_budget_limiter": _PROXY_MaxBudgetLimiter, "parallel_request_limiter": _PROXY_MaxParallelRequestsHandler_v3, "cache_control_check": _PROXY_CacheControlCheck, "responses_id_security": ResponsesIDSecurity, @@ -35,7 +33,7 @@ if os.getenv("LEGACY_MULTI_INSTANCE_RATE_LIMITING", "false").lower() == "true": def get_proxy_hook( - hook_name: Literal["max_budget_limiter", "managed_files", "parallel_request_limiter", "cache_control_check"] | str, + hook_name: Literal["managed_files", "parallel_request_limiter", "cache_control_check"] | str, ): """ Factory method to get a proxy hook instance by name diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py deleted file mode 100644 index eaf37b0bcf1..00000000000 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ /dev/null @@ -1,84 +0,0 @@ -from typing import Final - -from fastapi import HTTPException - -from litellm import verbose_logger -from litellm._logging import verbose_proxy_logger -from litellm.caching.caching import DualCache -from litellm.exceptions import RateLimitType -from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError -from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit - - -class _PROXY_MaxBudgetLimiter(CustomLogger): - # Class variables or attributes - def __init__(self): - pass - - async def async_pre_call_hook( - self, - user_api_key_dict: UserAPIKeyAuth, - cache: DualCache, - data: dict, - call_type: str, - ): - try: - verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook") - max_budget: Final = user_api_key_dict.user_max_budget - user_id: Final = user_api_key_dict.user_id - - if max_budget is None or user_id is None: - return - - from litellm.proxy.proxy_server import general_settings - - if ( - user_api_key_dict.team_id is not None - and general_settings.get("apply_user_budget_to_team_keys") is not True - ): - return - - # The reservation path admits at the strict-`<` boundary and - # atomically pre-fills the same counter we'd read here. Re-checking - # with `>=` would reject a request the reservation already admitted - # when the reservation fills the counter to exactly max_budget. - # Imported lazily to avoid a circular import via proxy.utils. - from litellm.proxy.spend_tracking.budget_reservation import ( - get_reserved_counter_keys, - ) - - user_counter_key: Final = f"spend:user:{user_id}" - if user_counter_key in get_reserved_counter_keys(user_api_key_dict.budget_reservation): - return - - from litellm.proxy.proxy_server import get_current_spend - - curr_spend: Final = await get_current_spend( - counter_key=user_counter_key, - fallback_spend=user_api_key_dict.user_spend or 0.0, - ) - - verbose_proxy_logger.debug( - "MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f", - user_id, - curr_spend, - max_budget, - ) - - # CHECK IF REQUEST ALLOWED - if curr_spend >= max_budget: - resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None) - raise ProxyRateLimitError( - detail="Max budget limit reached.", - rate_limit_type=RateLimitType.BUDGET, - model=resolved_model, - llm_provider=llm_provider, - ) - except HTTPException as e: - raise e - except Exception as e: - verbose_logger.exception( - "litellm.proxy.hooks.max_budget_limiter.py::async_pre_call_hook(): Exception occured - %s", e - ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 99cd78e0b58..18e58861266 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -164,7 +164,6 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai ) from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck -from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter from litellm.proxy.hooks.parallel_request_limiter import ( _PROXY_MaxParallelRequestsHandler, ) @@ -982,7 +981,6 @@ class ProxyLogging: dual_cache=DualCache(default_in_memory_ttl=1) # ping redis cache every 1s ) self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache) - self.max_budget_limiter = _PROXY_MaxBudgetLimiter() self.cache_control_check = _PROXY_CacheControlCheck() self.alerting: list[str] | None = None self.alerting_threshold: float = 300 # default to 5 min. threshold @@ -3580,7 +3578,7 @@ class ProxyLogging: caps: Final = ProxyLogging._callback_capabilities() post_call_pipelines: Final = _streamable_post_call_pipelines(request_data, user_api_key_dict) # Fast path: no real overrides. Internal proxy CustomLogger callbacks - # (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default + # (e.g. _PROXY_CacheControlCheck, ManagedFiles) inherit the default # ``async for chunk: yield chunk`` body, so wrapping the iterator # through each of them adds N pass-through trampolines per chunk for # zero behavior change. Skip the chain entirely and stream through. diff --git a/tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md b/tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md index 7ff920a8d6d..a07bdf3d4e9 100644 --- a/tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md +++ b/tests/e2e/quota_management/budgets/BUDGET_TEST_COVERAGE_MATRIX.md @@ -21,7 +21,7 @@ on the shared lifecycle (every entity it creates is deleted on teardown). | Entity | Unit | Pre-existing live | This suite (live) | Status | |--------|------|-------------------|-------------------|--------| -| API key | `test_budget_reservation.py`, `test_max_budget_limiter.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** | +| API key | `test_budget_reservation.py` | `otel_tests` | `test_budget_enforcement_e2e::test_key_budget_blocks` | **covered** | | Team | `test_team_budget_limits.py` | `otel_tests` | (org test builds a team) | **covered** | | Internal user | auth unit tests | - | `test_internal_user_budget_blocks` | **covered (new)** | | Team member | `test_team_member_budget.py` | - | `test_team_member_budget_blocks` | **covered (new)** | diff --git a/tests/logging_callback_tests/test_unit_test_litellm_logging.py b/tests/logging_callback_tests/test_unit_test_litellm_logging.py index 42ba4ff35f1..7709a823610 100644 --- a/tests/logging_callback_tests/test_unit_test_litellm_logging.py +++ b/tests/logging_callback_tests/test_unit_test_litellm_logging.py @@ -8,8 +8,8 @@ from typing import Literal import pytest import litellm from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck +from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler from litellm._service_logger import ServiceLogging import asyncio @@ -58,11 +58,11 @@ def test_is_internal_litellm_proxy_callback(): """ Ensure we can determine if a callback is an internal litellm proxy callback - eg. `_PROXY_MaxBudgetLimiter`, `_PROXY_CacheControlCheck` + eg. `_PROXY_MaxIterationsHandler`, `_PROXY_CacheControlCheck` """ logging = setup_logging() - assert logging._is_internal_litellm_proxy_callback(_PROXY_MaxBudgetLimiter) == True + assert logging._is_internal_litellm_proxy_callback(_PROXY_MaxIterationsHandler) == True # Test non-internal callbacks def regular_callback(): @@ -95,7 +95,7 @@ def test_should_run_sync_callbacks_for_async_calls(): assert logging._should_run_sync_callbacks_for_async_calls() == True # Test with internal callback only - litellm.success_callback = [_PROXY_MaxBudgetLimiter] + litellm.success_callback = [_PROXY_MaxIterationsHandler] assert logging._should_run_sync_callbacks_for_async_calls() == False @@ -107,7 +107,7 @@ def test_remove_internal_litellm_callbacks(): callbacks = [ regular_callback, - _PROXY_MaxBudgetLimiter, + _PROXY_MaxIterationsHandler, _PROXY_CacheControlCheck, "string_callback", ] @@ -116,5 +116,5 @@ def test_remove_internal_litellm_callbacks(): assert len(filtered) == 2 # Should only keep regular_callback and string_callback assert regular_callback in filtered assert "string_callback" in filtered - assert _PROXY_MaxBudgetLimiter not in filtered + assert _PROXY_MaxIterationsHandler not in filtered assert _PROXY_CacheControlCheck not in filtered diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py deleted file mode 100644 index 71671966d1a..00000000000 --- a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py +++ /dev/null @@ -1,237 +0,0 @@ -""" -Unit tests for the personal-budget pre-call hook. - -The reservation path (added in PR #26845) atomically pre-fills the same -`spend:user:{user_id}` counter this hook reads, admitting at a strict-`<` -boundary. Re-checking with `>=` after reservation would reject requests the -reservation already admitted when the reservation fills the counter to -exactly `max_budget` (e.g. requests with no `max_tokens` cap fall back to -reserving the smallest remaining headroom). - -These tests pin the skip-when-reserved behavior and guard against drift. -""" - -from unittest.mock import AsyncMock, patch - -import pytest -from fastapi import HTTPException - -from litellm.caching.caching import DualCache -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter - - -def _make_user_api_key_auth( - user_id: str = "user-1", - user_max_budget: float = 10.0, - user_spend: float = 0.0, - team_id=None, - budget_reservation=None, -) -> UserAPIKeyAuth: - return UserAPIKeyAuth( - api_key="sk-test", - user_id=user_id, - user_max_budget=user_max_budget, - user_spend=user_spend, - team_id=team_id, - budget_reservation=budget_reservation, - ) - - -@pytest.mark.asyncio -async def test_under_budget_passes(): - handler = _PROXY_MaxBudgetLimiter() - user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0) - - with patch( - "litellm.proxy.proxy_server.get_current_spend", - new=AsyncMock(return_value=3.0), - ): - result = await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=DualCache(), - data={}, - call_type="completion", - ) - - assert result is None - - -@pytest.mark.asyncio -async def test_over_budget_rejects_without_reservation(): - handler = _PROXY_MaxBudgetLimiter() - user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0) - - with patch( - "litellm.proxy.proxy_server.get_current_spend", - new=AsyncMock(return_value=10.0), - ): - with pytest.raises(HTTPException) as exc_info: - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=DualCache(), - data={}, - call_type="completion", - ) - - assert exc_info.value.status_code == 429 - assert "Max budget limit reached." in exc_info.value.detail - - -@pytest.mark.asyncio -async def test_skips_when_user_counter_is_reserved(): - """ - Reservation atomically pre-fills `spend:user:{user_id}` and admits the - request. The legacy `>=` check must not double-enforce on the same - counter — that's what produced the boundary regression where a fresh - user with no `max_tokens` cap got 429'd on their first request. - """ - handler = _PROXY_MaxBudgetLimiter() - user_api_key_dict = _make_user_api_key_auth( - user_id="user-1", - user_max_budget=10.0, - budget_reservation={ - "reserved_cost": 10.0, - "entries": [ - { - "counter_key": "spend:user:user-1", - "entity_type": "User", - "entity_id": "user-1", - "reserved_cost": 10.0, - "applied_adjustment": 0.0, - } - ], - "finalized": False, - }, - ) - - # `get_current_spend` would return 10.0 here (counter pre-filled by the - # reservation). The hook must skip without reading it. - with patch( - "litellm.proxy.proxy_server.get_current_spend", - new=AsyncMock(return_value=10.0), - ) as mock_get_spend: - result = await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=DualCache(), - data={}, - call_type="completion", - ) - - assert result is None - mock_get_spend.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_does_not_skip_when_reservation_covers_a_different_counter(): - """ - A reservation that only covers e.g. `spend:team:{team_id}` (not the user - counter) must not exempt the user-budget check. - """ - handler = _PROXY_MaxBudgetLimiter() - user_api_key_dict = _make_user_api_key_auth( - user_id="user-1", - user_max_budget=10.0, - budget_reservation={ - "reserved_cost": 5.0, - "entries": [ - { - "counter_key": "spend:team:team-x", - "entity_type": "Team", - "entity_id": "team-x", - "reserved_cost": 5.0, - "applied_adjustment": 0.0, - } - ], - "finalized": False, - }, - ) - - with patch( - "litellm.proxy.proxy_server.get_current_spend", - new=AsyncMock(return_value=10.0), - ): - with pytest.raises(HTTPException) as exc_info: - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=DualCache(), - data={}, - call_type="completion", - ) - - assert exc_info.value.status_code == 429 - - -@pytest.mark.asyncio -async def test_team_keys_skip_personal_budget(): - handler = _PROXY_MaxBudgetLimiter() - user_api_key_dict = _make_user_api_key_auth( - user_max_budget=10.0, - team_id="team-1", - ) - - with patch( - "litellm.proxy.proxy_server.get_current_spend", - new=AsyncMock(return_value=999.0), - ) as mock_get_spend: - result = await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=DualCache(), - data={}, - call_type="completion", - ) - - assert result is None - mock_get_spend.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_team_keys_enforce_personal_budget_when_flag_enabled(): - """This hook is the third personal-budget gate alongside common_checks and the - reservation path, so apply_user_budget_to_team_keys has to reach it too or an - opted-in deployment enforces in two places out of three.""" - handler = _PROXY_MaxBudgetLimiter() - user_api_key_dict = _make_user_api_key_auth( - user_max_budget=10.0, - team_id="team-1", - ) - - with patch.dict( - "litellm.proxy.proxy_server.general_settings", - {"apply_user_budget_to_team_keys": True}, - ), patch( - "litellm.proxy.proxy_server.get_current_spend", - new=AsyncMock(return_value=999.0), - ): - with pytest.raises(HTTPException) as exc_info: - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=DualCache(), - data={}, - call_type="completion", - ) - - assert exc_info.value.status_code == 429 - - -@pytest.mark.asyncio -async def test_no_max_budget_passes(): - handler = _PROXY_MaxBudgetLimiter() - user_api_key_dict = UserAPIKeyAuth( - api_key="sk-test", - user_id="user-1", - ) - - with patch( - "litellm.proxy.proxy_server.get_current_spend", - new=AsyncMock(return_value=999.0), - ) as mock_get_spend: - result = await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=DualCache(), - data={}, - call_type="completion", - ) - - assert result is None - mock_get_spend.assert_not_awaited() diff --git a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py index ec680317980..49bbd498cb9 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py @@ -6,7 +6,7 @@ Background ---------- The proxy's internal rate-limit hooks (parallel_request_limiter, parallel_request_limiter_v3, dynamic_rate_limiter, dynamic_rate_limiter_v3, -batch_rate_limiter, max_budget_limiter, max_iterations_limiter, +batch_rate_limiter, max_iterations_limiter, max_budget_per_session_limiter) all fire from ``async_pre_call_hook`` — *before* :func:`litellm.get_llm_provider` runs anywhere else in the request lifecycle. @@ -50,7 +50,6 @@ from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHand from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( _PROXY_DynamicRateLimitHandlerV3, ) -from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter from litellm.proxy.hooks.max_budget_per_session_limiter import ( _PROXY_MaxBudgetPerSessionHandler, ) @@ -830,64 +829,6 @@ async def test_batch_rate_limiter_unknown_model_falls_back(): assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK -# --------------------------------------------------------------------------- -# max_budget_limiter -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_max_budget_limiter_populates_provider(): - handler = _PROXY_MaxBudgetLimiter() - user_api_key_dict = UserAPIKeyAuth( - api_key="sk-budget", - user_id="user-1", - user_max_budget=10.0, - ) - - with patch( - "litellm.proxy.proxy_server.get_current_spend", - new=AsyncMock(return_value=10.0), - ): - with pytest.raises(HTTPException) as exc_info: - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=DualCache(), - data={"model": "gpt-4o-mini"}, - call_type="completion", - ) - - exc = exc_info.value - assert exc.status_code == 429 - assert isinstance(exc, RateLimitError) - assert exc.llm_provider == "openai" - assert exc.model == "gpt-4o-mini" - - -@pytest.mark.asyncio -async def test_max_budget_limiter_no_model_falls_back(): - handler = _PROXY_MaxBudgetLimiter() - user_api_key_dict = UserAPIKeyAuth( - api_key="sk-budget", - user_id="user-1", - user_max_budget=10.0, - ) - - with patch( - "litellm.proxy.proxy_server.get_current_spend", - new=AsyncMock(return_value=10.0), - ): - with pytest.raises(HTTPException) as exc_info: - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=DualCache(), - data={}, - call_type="completion", - ) - - assert exc_info.value.llm_provider == PROXY_LLM_PROVIDER_FALLBACK - assert exc_info.value.model == "" - - # --------------------------------------------------------------------------- # max_iterations_limiter # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index 2d9c1bd8b46..dd3914e3ad5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -1397,7 +1397,7 @@ def test_get_config_callbacks_excludes_internal_runtime_callbacks(client, auth_a from litellm.integrations.s3_v2 import S3Logger from litellm.integrations.sqs import SQSLogger from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import VectorStorePreCallHook - from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter + from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck from litellm.router import Router class _InventoryTestGuardrail(CustomGuardrail): @@ -1425,7 +1425,7 @@ def test_get_config_callbacks_excludes_internal_runtime_callbacks(client, auth_a litellm, "callbacks", [ - _PROXY_MaxBudgetLimiter(), + _PROXY_CacheControlCheck(), _PROXY_LiteLLMManagedFiles(internal_usage_cache=MagicMock(), prisma_client=MagicMock()), ServiceLogging(), VectorStorePreCallHook(), diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 28ff4571b44..a3ff7f7447e 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -73,7 +73,7 @@ async def test_post_call_response_headers_hook_returns_early_without_callbacks( def test_callback_capabilities_skips_default_custom_logger(monkeypatch): """ - Internal proxy hooks (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit + Internal proxy hooks (e.g. _PROXY_CacheControlCheck, ManagedFiles) inherit the default ``async_post_call_streaming_iterator_hook`` body. The capability scanner must NOT report them as iterator overrides — wrapping the chunk stream through every no-op layer was responsible for ~10x diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py index a97dcb41e44..40cb3f10d34 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py @@ -220,7 +220,7 @@ def test_add_proxy_hooks_registers_callbacks(proxy_logging, monkeypatch): what gets registered. Verifies that the resulting instances land in ``proxy_logging.proxy_hook_mapping`` keyed by hook name. """ - hook_keys = ["cache_control_check", "max_budget_limiter"] + hook_keys = ["cache_control_check", "max_iterations_limiter"] registered: List[Any] = [] from litellm.proxy import utils as utils_mod @@ -362,22 +362,22 @@ def test_add_proxy_hooks_unknown_hook_raises(proxy_logging, monkeypatch): def test_get_proxy_hook_returns_registered_instance(proxy_logging): s_cache = MagicMock() - s_budget = MagicMock() + s_iterations = MagicMock() s_parallel = MagicMock() proxy_logging.proxy_hook_mapping = { "cache_control_check": s_cache, - "max_budget_limiter": s_budget, + "max_iterations_limiter": s_iterations, "max_parallel_request_limiter": s_parallel, } snapshot = { "cache_control_check": proxy_logging.get_proxy_hook("cache_control_check") is s_cache, - "max_budget_limiter": proxy_logging.get_proxy_hook("max_budget_limiter") is s_budget, + "max_iterations_limiter": proxy_logging.get_proxy_hook("max_iterations_limiter") is s_iterations, "max_parallel_request_limiter": proxy_logging.get_proxy_hook("max_parallel_request_limiter") is s_parallel, "unknown_returns_none": proxy_logging.get_proxy_hook("unknown") is None, } assert snapshot == { "cache_control_check": True, - "max_budget_limiter": True, + "max_iterations_limiter": True, "max_parallel_request_limiter": True, "unknown_returns_none": True, } diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index af89c424f8b..6e5cb7fcae3 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -3,7 +3,7 @@ from __future__ import annotations from typing import Any, Dict -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException @@ -400,6 +400,37 @@ def test_has_pre_call_guardrails_counts_a_content_enforcer(proxy_logging, monkey assert proxy_logging.has_pre_call_guardrails({}) is True +@pytest.mark.asyncio +async def test_registered_hooks_do_not_enforce_user_budget(proxy_logging, monkeypatch): + """ + Personal budget is auth's job (`_user_max_budget_check`), which exempts + zero-cost models. A hook re-checking the same counter without that + exemption is what 429'd free models once a user was over budget. + """ + monkeypatch.setattr(litellm, "callbacks", []) + with patch("litellm.proxy.proxy_server.prisma_client", None): + proxy_logging._add_proxy_hooks(llm_router=None) + ProxyLogging._callback_capabilities_cache.clear() + + over_budget_user = UserAPIKeyAuth( + api_key="sk-personal", + user_id="user-over-budget", + user_max_budget=1.0, + user_spend=5.0, + team_id=None, + ) + data = {"model": "free-model", "messages": [{"role": "user", "content": "hi"}]} + + with patch("litellm.proxy.proxy_server.get_current_spend", new=AsyncMock(return_value=5.0)): + out = await proxy_logging.pre_call_hook( + user_api_key_dict=over_budget_user, + data=data, + call_type="completion", + ) + + assert out == data + + def test_every_pre_call_customlogger_is_deliberately_classified(): """ A ledger, so a new hook cannot land unclassified. @@ -415,7 +446,6 @@ def test_every_pre_call_customlogger_is_deliberately_classified(): "_ENTERPRISE_BlockedUserList", } counts_or_shapes_the_request = { - "_PROXY_MaxBudgetLimiter", "_PROXY_MaxParallelRequestsHandler_v3", "_PROXY_MaxIterationsHandler", "_PROXY_MaxBudgetPerSessionHandler", diff --git a/tests/test_litellm/test_rate_limit_error_unification.py b/tests/test_litellm/test_rate_limit_error_unification.py index 99e9981857c..8241b29aff1 100644 --- a/tests/test_litellm/test_rate_limit_error_unification.py +++ b/tests/test_litellm/test_rate_limit_error_unification.py @@ -221,28 +221,6 @@ class TestProxyHookCategoryWiring: """End-to-end check that every proxy-side rate limiter raises the unified class with a sensible category, not a bare HTTPException.""" - def test_max_budget_limiter_raises_proxy_rate_limit_error(self): - from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter - - limiter = _PROXY_MaxBudgetLimiter() - # The simplest deterministic path: directly raise from the conditional - # branch by calling into the helper's exception construction. We - # round-trip through the public class to assert the shape. - with pytest.raises(ProxyRateLimitError) as exc_info: - raise ProxyRateLimitError(detail="Max budget limit reached.") - assert exc_info.value.status_code == 429 - assert exc_info.value.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT - # And it's also a RateLimitError + HTTPException (the unification). - assert isinstance(exc_info.value, RateLimitError) - assert isinstance(exc_info.value, HTTPException) - # Static check that the limiter's module imports the unified class so - # the source of truth is wired correctly. - from litellm.proxy.hooks import max_budget_limiter - - assert hasattr(max_budget_limiter, "ProxyRateLimitError") - assert max_budget_limiter.ProxyRateLimitError is ProxyRateLimitError - del limiter # silence unused-var - @pytest.mark.parametrize( "module_path", [ @@ -251,7 +229,6 @@ class TestProxyHookCategoryWiring: "litellm.proxy.hooks.dynamic_rate_limiter", "litellm.proxy.hooks.dynamic_rate_limiter_v3", "litellm.proxy.hooks.batch_rate_limiter", - "litellm.proxy.hooks.max_budget_limiter", "litellm.proxy.hooks.max_budget_per_session_limiter", "litellm.proxy.hooks.max_iterations_limiter", ], @@ -542,44 +519,6 @@ class TestProxyHooksActuallyRaiseProxyRateLimitError: assert isinstance(e, RateLimitError) assert isinstance(e, HTTPException) - @pytest.mark.asyncio - async def test_max_budget_limiter_raises_proxy_rate_limit_error(self): - """ - Drive `_PROXY_MaxBudgetLimiter` past the user budget and assert it - raises the unified class. Mocks `get_current_spend` so we don't need - the proxy DB. - """ - from unittest.mock import patch - - from litellm.caching.caching import DualCache - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.hooks.max_budget_limiter import ( - _PROXY_MaxBudgetLimiter, - ) - - handler = _PROXY_MaxBudgetLimiter() - user_api_key_dict = UserAPIKeyAuth( - api_key="sk-test-budget", - user_id="user-budget-1", - user_max_budget=1.0, - user_spend=2.0, - ) - with patch( - "litellm.proxy.proxy_server.get_current_spend", - return_value=5.0, - ): - with pytest.raises(ProxyRateLimitError) as exc_info: - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=DualCache(), - data={}, - call_type="completion", - ) - e = exc_info.value - assert e.status_code == 429 - assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT - assert "max budget" in str(e.detail).lower() - @pytest.mark.asyncio async def test_dynamic_rate_limiter_v1_raises_proxy_rate_limit_error(self): """ @@ -1156,14 +1095,6 @@ class TestProxyHooksWireTypeCorrectly: max-iterations) without grepping the error message. """ - def test_max_budget_limiter_emits_budget_type(self): - e = ProxyRateLimitError( - detail="Max budget limit reached.", - rate_limit_type=RateLimitType.BUDGET, - ) - assert e.category == "litellm_rate_limit" - assert e.rate_limit_type == "budget" - def test_max_iterations_limiter_emits_max_iterations_type(self): e = ProxyRateLimitError( detail="Max iterations exceeded for session abc.", From dae16264c1b2b0c39face0014aff8c1a0028e606 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:21:33 +0000 Subject: [PATCH 151/187] test(auth): cover over-budget user on zero-cost vs paid model in common_checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/auth/test_auth_checks.py | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26ae28a57d2..6b8260e7f08 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5605,6 +5605,65 @@ async def test_common_checks_personal_user_budget_blocks_in_gather(): assert "User=u1" in str(over.value) +async def _common_checks_for_over_budget_personal_key(*, model: str) -> bool: + from litellm import Router + from litellm.proxy.auth.auth_checks import _is_model_cost_zero, common_checks + + llm_router: Final = Router( + model_list=[ + { + "model_name": "free-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}, + }, + { + "model_name": "paid-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + }, + ] + ) + user: Final = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=1.0) + token: Final = UserAPIKeyAuth(token="k1", user_id="u1") + + async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): + return 5.0 if counter_key == "spend:user:u1" else 0.0 + + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.get_current_spend", _spend_by_counter), + ): + result: Final = await common_checks( + request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=llm_router, + proxy_logging_obj=proxy_logging_obj, + valid_token=token, + request=MagicMock(spec=Request), + skip_budget_checks=_is_model_cost_zero(model=model, llm_router=llm_router), + ) + await asyncio.sleep(0) + return result + + +@pytest.mark.asyncio +async def test_common_checks_over_budget_user_can_still_call_zero_cost_model(): + """LIT-7464: an exhausted personal budget must not block a model priced at 0/0, + while the same user is still rejected on a priced model.""" + assert await _common_checks_for_over_budget_personal_key(model="free-model") is True + + with pytest.raises(litellm.BudgetExceededError) as over: + await _common_checks_for_over_budget_personal_key(model="paid-model") + assert "ExceededBudget: User=u1" in str(over.value) + + async def _run_internal_user_budget_alert( *, spend: float, From ebae692a0ddc854bebbec41983ad1fe5ba070b66 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 15:45:50 -0700 Subject: [PATCH 152/187] refactor(responses): drop the api_base cast and mark the header merge mutable-ok --- litellm/llms/azure_ai/responses/transformation.py | 6 +++++- litellm/responses/main.py | 9 +++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure_ai/responses/transformation.py b/litellm/llms/azure_ai/responses/transformation.py index b61c856f733..66a284c821d 100644 --- a/litellm/llms/azure_ai/responses/transformation.py +++ b/litellm/llms/azure_ai/responses/transformation.py @@ -34,7 +34,11 @@ class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig): litellm_params=params.model_dump(), api_key_header=api_key_header_for_base(AzureFoundryModelInfo.get_api_base(params.api_base)), ) - return {**headers, **auth_headers, "Content-Type": "application/json"} + return { # mutable-ok: the handler updates the returned headers in place per the dict contract + **headers, + **auth_headers, + "Content-Type": "application/json", + } def supports_native_websocket(self) -> bool: return False diff --git a/litellm/responses/main.py b/litellm/responses/main.py index a138d6f8eb3..63bee9f6d99 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -492,6 +492,11 @@ def _resolve_responses_api_provider_config( return OpenAILikeResponsesConfig() +def _api_base_kwarg(kwargs: Mapping[str, object]) -> str | None: + api_base: Final = kwargs.get("api_base") + return api_base if isinstance(api_base, str) else None + + def _will_bridge_to_chat_completions( model: str, custom_llm_provider: str | None, @@ -622,7 +627,7 @@ async def aresponses( custom_llm_provider, bool(kwargs.get("use_chat_completions_api")), kwargs.get("model_info"), - cast(str | None, kwargs.get("api_base")), + _api_base_kwarg(kwargs), ), ): ( @@ -792,7 +797,7 @@ def _apply_prompt_management_to_responses_call( custom_llm_provider, use_chat_completions_api, kwargs.get("model_info"), - cast(str | None, kwargs.get("api_base")), + _api_base_kwarg(kwargs), ), ): ( From 4b70696afafcc11793cba55423fa6e77b0a90b29 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 22:49:25 +0000 Subject: [PATCH 153/187] fix(streaming): estimate interrupted Anthropic stream usage from reasoning_content Interrupted Anthropic streams that die before message_delta were billed at the message_start placeholder (any value above 1 was trusted) or at 0 when the partial response was reasoning-only, because the token_counter fallback only looked at visible text. Reset the placeholder whenever no finish_reason or second usage event arrived, fold the already-counted reasoning tokens into the fallback estimate, and drop the stale completion_tokens_details so cost is computed from the recovered count Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 70 +++++++++++---- .../test_streaming_chunk_builder_cursor.py | 86 +++++++++++++++++-- 2 files changed, 132 insertions(+), 24 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 90698296142..a8b1f81702c 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -5,6 +5,7 @@ from itertools import groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast +from pydantic import BaseModel from typing_extensions import ReadOnly, Required from litellm._logging import verbose_logger @@ -148,6 +149,7 @@ class _ToolCallChunk(TypedDict): class _UsageBearingChunk(TypedDict, total=False): usage: Usage | None _hidden_params: Mapping[str, str] + choices: ReadOnly[Sequence[StreamingChoices | Mapping[str, object]]] class _UsageSummary(TypedDict): @@ -921,21 +923,22 @@ class ChunkProcessor: prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details) - completion_tokens = self._reset_anthropic_cursor_completion_tokens( + recovered_completion_tokens: Final = self._reset_anthropic_cursor_completion_tokens( chunks=chunks, completion_tokens=completion_tokens, completion_usage_updates=completion_usage_updates, ) + cursor_was_reset: Final = recovered_completion_tokens != completion_tokens return UsagePerChunk( prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, + completion_tokens=recovered_completion_tokens, cache_creation_input_tokens=cache_creation_input_tokens, cache_read_input_tokens=cache_read_input_tokens, server_tool_use=server_tool_use, web_search_requests=web_search_requests, google_maps_grounding_requests=google_maps_grounding_requests, - completion_tokens_details=completion_tokens_details, + completion_tokens_details=None if cursor_was_reset else completion_tokens_details, prompt_tokens_details=prompt_tokens_details, cost=cost, inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"), @@ -960,6 +963,32 @@ class ChunkProcessor: ] return values[-1] if values else None + @staticmethod + def _finish_reason_of_choice(choice: object) -> str | None: + match choice: + case StreamingChoices(finish_reason=reason) | Choices(finish_reason=reason): + return reason + case {"finish_reason": str() as reason}: + return reason + case _: + return None + + @staticmethod + def _chunk_choices(chunk: "_UsageBearingChunk | BaseModel") -> Sequence[object]: + if isinstance(chunk, dict): + return chunk.get("choices", ()) + if isinstance(chunk, (ModelResponse, ModelResponseStream)): + return chunk.choices + return () + + @staticmethod + def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool: + return any( + ChunkProcessor._finish_reason_of_choice(choice) is not None + for chunk in chunks + for choice in ChunkProcessor._chunk_choices(chunk) + ) + @staticmethod def _reset_anthropic_cursor_completion_tokens( chunks: Sequence["_UsageBearingChunk | ModelResponse"], @@ -970,18 +999,18 @@ class ChunkProcessor: See the ``completion_usage_updates`` comment in ``_calculate_usage_per_chunk``. The accumulated value is NOT a stale - cursor when either it is > 1 (definitely not a placeholder) or we saw - >= 2 completion-bearing usage events (positive evidence ``message_delta`` - arrived). Otherwise — the only completion update we ever saw was the - Anthropic ``message_start`` cursor (=1) — reset to 0 so - ``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates - from the actually-received completion text instead of trusting the - placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the - heuristic (which encodes Anthropic's specific message_start SSE shape) - does not silently affect other providers that may legitimately report - ``completion_tokens=1`` from a single usage event. + cursor when we saw >= 2 completion-bearing usage events or any chunk + carried a ``finish_reason`` (positive evidence ``message_delta`` + arrived). Otherwise the only completion update we ever saw was the + Anthropic ``message_start`` cursor, a small placeholder whose magnitude + varies per request (1 and 8 both observed live), so reset to 0 and let + ``calculate_usage()``'s ``or token_counter(...)`` fallback estimate from + the actually-received text and reasoning instead. Gated on + ``custom_llm_provider == "anthropic"`` so the heuristic (which encodes + Anthropic's specific message_start SSE shape) does not silently affect + other providers that legitimately report usage from a single event. """ - saw_non_cursor_completion: Final = completion_tokens > 1 or completion_usage_updates >= 2 + saw_non_cursor_completion: Final = completion_usage_updates >= 2 or ChunkProcessor._saw_finish_reason(chunks) if saw_non_cursor_completion: return completion_tokens @@ -995,7 +1024,7 @@ class ChunkProcessor: if isinstance(hp, dict): custom_llm_provider = hp.get("custom_llm_provider") - if custom_llm_provider == "anthropic" and completion_tokens == 1: + if custom_llm_provider == "anthropic": return 0 return completion_tokens @@ -1039,10 +1068,13 @@ class ChunkProcessor: returned_usage.prompt_tokens = 0 returned_usage.completion_tokens = ( completion_tokens - or token_counter( - model=model, - text=completion_output, - count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages + or ( + token_counter( + model=model, + text=completion_output, + count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages + ) + + (reasoning_tokens or 0) ) ) returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py index 3d9971034ae..f4dbb28533f 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py @@ -22,9 +22,10 @@ text-based fallback to estimate from the real completion text. import pytest - +import litellm from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor from litellm.types.utils import ( + CompletionTokensDetailsWrapper, Delta, ModelResponseStream, StreamingChoices, @@ -35,6 +36,7 @@ from litellm.types.utils import ( def _make_chunk( *, content: str = "", + reasoning_content: str | None = None, usage: Usage = None, finish_reason: str = None, custom_llm_provider: str = "anthropic", @@ -48,7 +50,7 @@ def _make_chunk( StreamingChoices( finish_reason=finish_reason, index=0, - delta=Delta(content=content, role="assistant"), + delta=Delta(content=content, role="assistant", reasoning_content=reasoning_content), ) ], usage=usage, @@ -253,6 +255,79 @@ class TestAnthropicCursorBug: "Reset to 0 forces token_counter fallback." ) + @pytest.mark.parametrize("placeholder", [1, 3, 8]) + def test_interrupted_reasoning_only_stream_estimates_from_reasoning(self, placeholder: int): + """ + message_start placeholders are not always 1 (live Anthropic streams + have been observed sending 1 and 8 for the same prompt), and a thinking + model cut off before message_delta has streamed only reasoning_content. + The recovered usage, including the completion_tokens_details the cost + calculator bills from, must come from that reasoning rather than from + the placeholder. + """ + message_start = _make_chunk( + usage=Usage( + prompt_tokens=100, + completion_tokens=placeholder, + total_tokens=100 + placeholder, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=placeholder), + ) + ) + reasoning_text = "Let me work through the scheduling constraints step by step. " * 40 + reasoning_chunks = [ + _make_chunk(reasoning_content=reasoning_text[i : i + 50]) for i in range(0, len(reasoning_text), 50) + ] + + response = litellm.stream_chunk_builder( + chunks=[message_start, *reasoning_chunks], + messages=[{"role": "user", "content": "Plan the schedule."}], + ) + + assert response.choices[0].message.reasoning_content == reasoning_text + reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens + assert reasoning_tokens > placeholder + assert response.usage.completion_tokens == reasoning_tokens, ( + f"Expected completion_tokens to be the reasoning estimate, got " + f"completion_tokens={response.usage.completion_tokens} reasoning_tokens={reasoning_tokens}" + ) + assert response.usage.total_tokens == response.usage.prompt_tokens + reasoning_tokens + details = response.usage.completion_tokens_details + assert (details.text_tokens or 0) + details.reasoning_tokens == response.usage.completion_tokens + + def test_fallback_counts_reasoning_and_text_together(self): + """ + With no usable provider count, the estimate covers everything the + provider generated: reasoning_content plus visible text, not text alone. + """ + reasoning = "First I should check whether the input is sorted. " * 10 + text = "The list is already sorted, so no work is needed." + chunks = [_make_chunk(reasoning_content=reasoning), _make_chunk(content=text)] + + response = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "Sort it."}]) + + text_only = litellm.token_counter(model="claude-sonnet-4-6", text=text, count_response_tokens=True) + reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens + assert reasoning_tokens > 0 + assert response.usage.completion_tokens == text_only + reasoning_tokens + + def test_lone_usage_event_with_finish_reason_is_trusted(self): + """ + Guardrails rebuild responses from the chunks yielded to the client, + which excludes the un-yielded message_start. A finished stream then has + exactly one usage event (message_delta) and it must be kept as-is. + """ + chunks = [ + _make_chunk(content="Yes, "), + _make_chunk(content="that works."), + _make_chunk( + usage=Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25), + finish_reason="stop", + ), + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 5 + class TestProviderGuard: """Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic @@ -297,11 +372,12 @@ class TestNonAnthropicStreamingIntact: """Make sure providers without cursor pattern still work.""" def test_completion_tokens_above_one_never_resets(self): - """Any chunk reporting completion_tokens > 1 sets saw_non_cursor - and prevents the reset.""" + """A non-Anthropic provider reporting completion_tokens > 1 from a + single usage event keeps that value.""" chunks = [ _make_chunk( - usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + custom_llm_provider="openai", ), ] processor = ChunkProcessor(chunks=chunks, messages=[]) From c37d0a2e66a25a114e0af8f5d401e8b0ade1f545 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 15:51:04 -0700 Subject: [PATCH 154/187] fix(proxy): read general_settings without a cast and test the /team/update gate by behavior only --- .../proxy/management_endpoints/team_endpoints.py | 4 ++-- tests/test_litellm/proxy/auth/test_route_checks.py | 13 ++++--------- 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f0ec3f975d5..72367b2bced 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -504,9 +504,9 @@ _GENERAL_SETTINGS: Final = TypeAdapter(dict[str, object]) def _general_settings() -> Mapping[str, object]: - from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import general_settings - return _GENERAL_SETTINGS.validate_python(cast(object, proxy_server.general_settings)) + return _GENERAL_SETTINGS.validate_python(general_settings) def _caller_edit_access(role: TeamAccessRole | None, general_settings: Mapping[str, object]) -> TeamEditAccess: diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 2a6ba56a242..3211e85ff97 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -2897,10 +2897,6 @@ def test_team_update_gate_admits_internal_user_without_org_context(): caller and update_team resolves proxy, org or team admin itself, then filters team admins through the team_admin_editable_team_fields setting. Before that the gate 401'd every team admin, which left the handler's team-admin branch unreachable.""" - from litellm.proxy._types import LiteLLMRoutes - - assert "/team/update" in LiteLLMRoutes.self_managed_routes.value - user_obj = LiteLLM_UserTable( user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value, @@ -3001,11 +2997,10 @@ async def test_add_team_org_context_noop_for_static_team_route(): assert out == body -def test_patch_team_route_stays_out_of_self_managed_routes(): - """Unlike POST /team/update, PATCH /team/{team_id} cannot be self-managed: its - template also matches /team/new (the collision footgun), so it stays reachable by - org admins (org_admin_allowed_routes) and proxy admins only, never by regular - internal users or through the role-agnostic self_managed_routes.""" +def test_patch_team_route_has_same_reach_as_team_update(): + """/team/{team_id} is reachable by org admins (in org_admin_allowed_routes) but + NOT by regular internal users or the role-agnostic self_managed_routes — the + latter would open /team/new (the collision footgun) to any authenticated user.""" from litellm.proxy._types import LiteLLMRoutes assert RouteChecks.check_route_access( From eda81fff595f992bfae6471ef474eec896992fdc Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 15 Sep 2026 16:46:52 -0700 Subject: [PATCH 155/187] feat(ui): shared URL-state layer for tables and tabs Add useUrlTableState (search, sort, page, page size and filter_ in the query string via one nuqs useQueryStates call, with keyPrefix and urlKeys for routes that host two tables or need legacy key names) and useUrlTab (validated ?tab= param with a role-aware fallback). Migrate the Virtual Keys table onto useUrlTableState with a byte-identical URL contract and bind the Playground tab strip to ?tab=. DataTable gains controlled columnVisibility/onColumnVisibilityChange plus usePersistedColumnVisibility (localStorage per table id), and an isError prop that keeps the server page clamp from rewriting a deep-linked ?page= after a failed fetch. Virtual Keys uses both. The expired-session redirect in handleError now keeps the query string and hash so the return URL captured on re-login restores the filtered view instead of the bare list. Delete useTabRouting and tabRoutes, the pathname tab router left over from the reverted path-per-tab attempt (#34327, reverted in #34629); tab persistence has to be a query param on the static export. --- .../(dashboard)/hooks/useTabRouting.test.tsx | 82 ----- .../app/(dashboard)/hooks/useTabRouting.ts | 38 -- .../app/(dashboard)/playground/page.test.tsx | 53 ++- .../src/app/(dashboard)/playground/page.tsx | 10 +- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 152 +++++++- .../VirtualKeysPage/VirtualKeysTable.tsx | 185 ++++------ .../src/components/networking.test.ts | 38 +- .../src/components/networking.tsx | 2 +- .../shared/DataTable/DataTable.test-d.tsx | 27 +- .../shared/DataTable/DataTable.test.tsx | 98 +++++- .../components/shared/DataTable/DataTable.tsx | 15 +- .../src/components/shared/DataTable/index.ts | 2 + .../src/components/shared/DataTable/types.ts | 19 + .../usePersistedColumnVisibility.test.tsx | 98 ++++++ .../DataTable/usePersistedColumnVisibility.ts | 54 +++ .../DataTable/useUrlTableState.test.tsx | 325 ++++++++++++++++++ .../shared/DataTable/useUrlTableState.ts | 232 +++++++++++++ .../src/hooks/useUrlTab.test.tsx | 100 ++++++ ui/litellm-dashboard/src/hooks/useUrlTab.ts | 12 + .../src/utils/tabRoutes.test.ts | 47 --- ui/litellm-dashboard/src/utils/tabRoutes.ts | 26 -- 21 files changed, 1260 insertions(+), 355 deletions(-) delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts create mode 100644 ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts create mode 100644 ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx create mode 100644 ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts create mode 100644 ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx create mode 100644 ui/litellm-dashboard/src/hooks/useUrlTab.ts delete mode 100644 ui/litellm-dashboard/src/utils/tabRoutes.test.ts delete mode 100644 ui/litellm-dashboard/src/utils/tabRoutes.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx deleted file mode 100644 index 24900bae798..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.test.tsx +++ /dev/null @@ -1,82 +0,0 @@ -/* @vitest-environment jsdom */ -import { renderHook } from "@testing-library/react"; -import { beforeEach, describe, expect, it, vi } from "vitest"; - -const { mockPush, navState } = vi.hoisted(() => ({ - mockPush: vi.fn(), - navState: { pathname: "/logs" }, -})); -vi.mock("next/navigation", () => ({ - usePathname: () => navState.pathname, - useRouter: () => ({ push: mockPush }), -})); - -vi.mock("@/components/networking", () => ({ serverRootPath: "" })); - -import { createTabRoutes } from "@/utils/tabRoutes"; -import { useTabRouting } from "./useTabRouting"; - -const routes = createTabRoutes("logs", ["audit", "deleted-keys", "deleted-teams"] as const); - -const render = (ready = true) => { - const config = { - routes, - baseTabKey: "request-logs", - visibleKeys: ["audit", "deleted-keys", "deleted-teams"], - ready, - }; - return renderHook(() => useTabRouting(config)); -}; - -describe("useTabRouting", () => { - beforeEach(() => { - navState.pathname = "/logs"; - mockPush.mockClear(); - }); - - it("maps the base path to the base tab key", () => { - const { result } = render(); - expect(result.current.activeSlug).toBe(""); - expect(result.current.activeKey).toBe("request-logs"); - }); - - it("uses the slug itself as the active key for a known nested tab", () => { - navState.pathname = "/ui/logs/audit"; - const { result } = render(); - expect(result.current.activeKey).toBe("audit"); - }); - - it("falls back to the base tab key for an unknown slug", () => { - navState.pathname = "/ui/logs/bogus"; - const { result } = render(); - expect(result.current.activeKey).toBe("request-logs"); - }); - - it("redirects an unknown slug to the base href once ready", () => { - const replaceMock = vi.fn(); - const originalLocation = window.location; - Object.defineProperty(window, "location", { configurable: true, value: { replace: replaceMock } }); - navState.pathname = "/ui/logs/bogus"; - render(true); - expect(replaceMock).toHaveBeenCalledWith("/ui/logs/"); - Object.defineProperty(window, "location", { configurable: true, value: originalLocation }); - }); - - it("does not redirect while not ready (role/creds still loading)", () => { - const replaceMock = vi.fn(); - const originalLocation = window.location; - Object.defineProperty(window, "location", { configurable: true, value: { replace: replaceMock } }); - navState.pathname = "/ui/logs/bogus"; - render(false); - expect(replaceMock).not.toHaveBeenCalled(); - Object.defineProperty(window, "location", { configurable: true, value: originalLocation }); - }); - - it("pushes the tab href on change, mapping the base key back to the empty slug", () => { - const { result } = render(); - result.current.onTabChange("audit"); - expect(mockPush).toHaveBeenCalledWith("/ui/logs/audit/"); - result.current.onTabChange("request-logs"); - expect(mockPush).toHaveBeenCalledWith("/ui/logs/"); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts deleted file mode 100644 index c17d71b4855..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useTabRouting.ts +++ /dev/null @@ -1,38 +0,0 @@ -import { useEffect } from "react"; -import { usePathname, useRouter } from "next/navigation"; -import type { TabRoutes } from "@/utils/tabRoutes"; - -interface UseTabRoutingArgs { - routes: Pick, "tabHref" | "slugFromPathname">; - baseTabKey: string; - visibleKeys: readonly string[]; - ready?: boolean; -} - -interface TabRoutingState { - activeSlug: string; - activeKey: string; - onTabChange: (key: string) => void; -} - -export function useTabRouting({ routes, baseTabKey, visibleKeys, ready = true }: UseTabRoutingArgs): TabRoutingState { - const { tabHref, slugFromPathname } = routes; - const pathname = usePathname(); - const router = useRouter(); - - const activeSlug = slugFromPathname(pathname); - const isKnownSlug = activeSlug === "" || visibleKeys.includes(activeSlug); - const activeKey = isKnownSlug ? activeSlug || baseTabKey : baseTabKey; - - useEffect(() => { - if (ready && activeSlug !== "" && !isKnownSlug) { - window.location.replace(tabHref("")); - } - }, [ready, activeSlug, isKnownSlug, tabHref]); - - const onTabChange = (key: string) => { - router.push(tabHref(key === baseTabKey ? "" : key)); - }; - - return { activeSlug, activeKey, onTabChange }; -} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx index 85e19d7d251..f12ebc0b831 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.test.tsx @@ -1,5 +1,8 @@ -import { render, screen } from "@testing-library/react"; +import { screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; import PlaygroundPage from "./page"; const authState = { userRole: "Admin" }; @@ -35,14 +38,17 @@ vi.mock("@/app/(dashboard)/playground/components/chat_ui/AgentBuilderView", () = default: () =>
, })); -describe("PlaygroundPage role guard", () => { - beforeEach(() => { - authState.userRole = "Admin"; - }); +const lastUrlUpdate = (onUrlUpdate: ReturnType>) => + onUrlUpdate.mock.calls.at(-1)?.[0]; +beforeEach(() => { + authState.userRole = "Admin"; +}); + +describe("PlaygroundPage role guard", () => { it.each(["Internal Viewer", "Admin Viewer"])("blocks the entire playground for %s", (role) => { authState.userRole = role; - render(); + renderWithProviders(); expect(screen.getByText("Access Denied")).toBeInTheDocument(); expect(screen.queryByRole("tab")).not.toBeInTheDocument(); @@ -54,10 +60,43 @@ describe("PlaygroundPage role guard", () => { it.each(["Admin", "Internal User", "Org Admin"])("renders the playground for %s", (role) => { authState.userRole = role; - render(); + renderWithProviders(); expect(screen.queryByText("Access Denied")).not.toBeInTheDocument(); expect(screen.getByRole("tab", { name: "Chat" })).toBeInTheDocument(); expect(screen.getByTestId("chat-ui")).toBeInTheDocument(); }); }); + +describe("PlaygroundPage ?tab= deep link", () => { + it("opens on Chat when the URL has no tab", () => { + renderWithProviders(); + + expect(screen.getByRole("tab", { name: "Chat" })).toHaveAttribute("aria-selected", "true"); + }); + + it("activates the tab named in ?tab=", () => { + renderWithProviders(, { searchParams: { tab: "compare" } }); + + expect(screen.getByRole("tab", { name: "Compare" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("tab", { name: "Chat" })).toHaveAttribute("aria-selected", "false"); + }); + + it("falls back to Chat when ?tab= is not a playground tab", () => { + renderWithProviders(, { searchParams: { tab: "settings" } }); + + expect(screen.getByRole("tab", { name: "Chat" })).toHaveAttribute("aria-selected", "true"); + }); + + it("clicking a tab writes ?tab= with history replace", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + + await user.click(screen.getByRole("tab", { name: "Compliance" })); + + expect(await screen.findByRole("tab", { name: "Compliance", selected: true })).toBeInTheDocument(); + await waitFor(() => expect(lastUrlUpdate(onUrlUpdate)?.searchParams.get("tab")).toBe("compliance")); + expect(lastUrlUpdate(onUrlUpdate)?.options.history).toBe("replace"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx index 78ca538d8b5..27a61415672 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx @@ -9,6 +9,9 @@ import { DeprecationBanner } from "@/components/DeprecationBanner"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { fetchProxySettings } from "@/utils/proxyUtils"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { useUrlTab } from "@/hooks/useUrlTab"; + +const PLAYGROUND_TABS = ["chat", "compare", "compliance", "agent-builder"] as const; interface ProxySettings { PROXY_BASE_URL?: string; @@ -18,6 +21,7 @@ interface ProxySettings { export default function PlaygroundPage() { const { accessToken, userRole, userId, disabledPersonalKeyCreation, token, isViewOnly } = useAuthorized(); const [proxySettings, setProxySettings] = useState(undefined); + const [activeTab, setActiveTab] = useUrlTab(PLAYGROUND_TABS, "chat"); useEffect(() => { const initializeProxySettings = async () => { @@ -48,7 +52,11 @@ export default function PlaygroundPage() { return (
- + Chat diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 6c742309eb7..4d963a2f603 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -4,7 +4,7 @@ import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import { vi, it, expect, beforeEach, describe, Mock, MockedFunction } from "vitest"; import { chooseSelectOption, renderWithProviders } from "../../../tests/test-utils"; import { VirtualKeysTable } from "./VirtualKeysTable"; -import { KEY_TABLE_SORT_FIELDS } from "./keyTableColumns"; +import { KEY_TABLE_HIDDEN_COLUMNS, KEY_TABLE_SORT_FIELDS } from "./keyTableColumns"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import { useKeyInfo } from "@/app/(dashboard)/hooks/keys/useKeyInfo"; import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; @@ -187,6 +187,7 @@ const lastHistoryMode = (onUrlUpdate: Mock) => onUrlUpdate. beforeEach(() => { vi.clearAllMocks(); + localStorage.clear(); mockUseKeys.mockReturnValue(keysResult([mockKey])); mockUseKeyInfo.mockReturnValue(keyInfoResult(undefined)); @@ -823,16 +824,27 @@ describe("table state lives in the URL so it survives leaving and returning to t }); it("restores the drawer filters from the URL on mount", async () => { - renderWithProviders(, { searchParams: { filter_team: "team-1", filter_user: "user-42" } }); + const searchParams = { + filter_team: "team-1", + filter_org: "org-1", + filter_user: "user-42", + filter_key_id: mockKey.token, + }; + const expectedKeyListOptions = { + teamID: "team-1", + organizationID: "org-1", + userID: "user-42", + keyHash: mockKey.token, + }; + renderWithProviders(, { searchParams }); await waitFor(() => { - expect(mockUseKeys).toHaveBeenLastCalledWith( - 1, - 50, - expect.objectContaining({ teamID: "team-1", userID: "user-42" }), - ); + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining(expectedKeyListOptions)); }); expect(screen.getByTestId("filter-chip-team_id")).toHaveTextContent("Test Team"); + expect(screen.getByTestId("filter-chip-org_id")).toHaveTextContent("Test Organization"); + expect(screen.getByTestId("filter-chip-user_id")).toHaveTextContent("user-42"); + expect(screen.getByTestId("filter-chip-key_hash")).toHaveTextContent(mockKey.token); }); it("restores the status filter from the URL and sends it to /key/list", async () => { @@ -853,6 +865,21 @@ describe("table state lives in the URL so it survives leaving and returning to t expect(screen.queryByTestId("filter-chip-status")).not.toBeInTheDocument(); }); + it("drops a hand-edited status from the URL when another filter chip is removed", async () => { + const onUrlUpdate = vi.fn(); + renderWithProviders(, { + searchParams: { filter_status: "bogus", filter_user: "user-42" }, + onUrlUpdate, + }); + + fireEvent.click(await screen.findByTestId("filter-chip-remove-user_id")); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_user")).toBeNull(); + }); + expect(lastSearchParam(onUrlUpdate, "filter_status")).toBeNull(); + }); + it("writes the search term to the URL", async () => { const onUrlUpdate = vi.fn(); renderWithProviders(, { onUrlUpdate }); @@ -896,6 +923,40 @@ describe("table state lives in the URL so it survives leaving and returning to t expect(screen.queryByTestId("filter-chip-user_id")).not.toBeInTheDocument(); }); + it("writes the Organization and Key ID drawer filters to the URL and clears them again", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + + openFilters(); + await chooseSelectOption(user, await screen.findByPlaceholderText(/Select an organization/), /Test Organization/); + fireEvent.change(screen.getByPlaceholderText(/Enter Key ID/), { target: { value: mockKey.token } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_org")).toBe("org-1"); + }); + expect(lastSearchParam(onUrlUpdate, "filter_key_id")).toBe(mockKey.token); + expect(lastSearchParam(onUrlUpdate, "filter_org_id")).toBeNull(); + expect(lastSearchParam(onUrlUpdate, "filter_key_hash")).toBeNull(); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ organizationID: "org-1", keyHash: mockKey.token }), + ); + }); + + fireEvent.click(screen.getByTestId("datatable-clear-filters")); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_org")).toBeNull(); + }); + expect(lastSearchParam(onUrlUpdate, "filter_key_id")).toBeNull(); + expect(screen.queryByTestId("filter-chip-org_id")).not.toBeInTheDocument(); + expect(screen.queryByTestId("filter-chip-key_hash")).not.toBeInTheDocument(); + }); + it("returns to page 1 when the search term changes", async () => { const onUrlUpdate = vi.fn(); renderWithProviders(, { searchParams: { page: "3" }, onUrlUpdate }); @@ -953,14 +1014,16 @@ describe("table state lives in the URL so it survives leaving and returning to t }); }); - it("falls back to the default sort when the URL names a column the table cannot sort by", async () => { - renderWithProviders(, { searchParams: { sort_by: "totally_unknown_field" } }); + it("falls back to the default sort column, keeping the URL's direction, when the table cannot sort by sort_by", async () => { + renderWithProviders(, { + searchParams: { sort_by: "totally_unknown_field", sort_order: "asc" }, + }); await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith( 1, 50, - expect.objectContaining({ sortBy: "created_at", sortOrder: "desc" }), + expect.objectContaining({ sortBy: "created_at", sortOrder: "asc" }), ); }); expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); @@ -1003,3 +1066,72 @@ describe("table state lives in the URL so it survives leaving and returning to t }); }); }); + +describe("column choices survive a reload", () => { + const STORAGE_KEY = "litellm_table_columns_virtual-keys"; + const storedColumns = () => JSON.parse(localStorage.getItem(STORAGE_KEY) ?? "null"); + + it("hides a column that was hidden on a previous visit while the default-hidden columns stay hidden", () => { + localStorage.setItem(STORAGE_KEY, JSON.stringify({ budget_reset_at: false })); + + renderWithProviders(); + + expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); + expect(screen.queryByText("Budget Reset")).not.toBeInTheDocument(); + expect(screen.queryByText("Created By")).not.toBeInTheDocument(); + }); + + it("writes a column toggled on through the Columns menu to storage and shows it again on the next mount", async () => { + const user = userEvent.setup(); + const { unmount } = renderWithProviders(); + expect(screen.queryByText("Created By")).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Columns" })); + await user.click(await screen.findByText("Created By")); + await user.keyboard("{Escape}"); + + expect(storedColumns()).toEqual({ ...KEY_TABLE_HIDDEN_COLUMNS, created_by: true }); + + unmount(); + renderWithProviders(); + + expect(screen.getByText("Created By")).toBeInTheDocument(); + }); +}); + +describe("a failed keys fetch does not rewrite the URL", () => { + const renderOnPage3OfMany = async () => { + mockUseKeys.mockReturnValue(keysResult([mockKey], { total_count: 200, total_pages: 4 })); + const onUrlUpdate = vi.fn(); + const view = renderWithProviders(, { searchParams: { page: "3" }, onUrlUpdate }); + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(3, 50, expect.anything()); + }); + return { ...view, onUrlUpdate }; + }; + + it("keeps ?page=3 when the keys query errors, instead of snapping to page 1 on the empty count", async () => { + const { rerender, onUrlUpdate } = await renderOnPage3OfMany(); + + mockUseKeys.mockReturnValue(keysResult([], {}, { data: undefined, isError: true })); + rerender(); + + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(mockUseKeys).toHaveBeenLastCalledWith(3, 50, expect.anything()); + expect(onUrlUpdate).not.toHaveBeenCalled(); + }); + + it("still snaps ?page=3 back to the first page when the keys query succeeds with no rows", async () => { + const { rerender, onUrlUpdate } = await renderOnPage3OfMany(); + + mockUseKeys.mockReturnValue(keysResult([])); + rerender(); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.anything()); + }); + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "page")).toBeNull(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 1f52bdd7335..39dd2cc5ab2 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -10,15 +10,18 @@ import { DataTableFilterDrawer, DataTableFilterField, DataTableToolbar, + usePersistedColumnVisibility, + useUrlTableState, + type UrlTableStateOptions, } from "@/components/shared/DataTable"; import { SearchSelect } from "@/components/shared/SearchSelect"; import { PageHeader } from "@/components/shared/PageHeader"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; -import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import { ColumnFiltersState, functionalUpdate, OnChangeFn } from "@tanstack/react-table"; import { KeyRound } from "lucide-react"; -import { createParser, parseAsInteger, parseAsString, parseAsStringLiteral, useQueryState, useQueryStates } from "nuqs"; +import { parseAsString, useQueryState } from "nuqs"; import React, { useCallback, useMemo, useState } from "react"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; @@ -56,44 +59,30 @@ const STATUS_FILTER_ITEMS = [ ...KEY_STATUS_VALUES.map((value) => ({ value, label: KEY_STATUS_LABELS[value] })), ]; -const isKeyStatusFilter = (value: string): value is KeyStatusFilter => - (KEY_STATUS_VALUES as readonly string[]).includes(value); +const isKeyStatusFilter = (value: unknown): value is KeyStatusFilter => + (KEY_STATUS_VALUES as readonly unknown[]).includes(value); -const DEFAULT_SORT_BY = "created_at"; -const DEFAULT_SORT_ORDER = "desc"; -const DEFAULT_PAGE_SIZE = 50; -const MAX_PAGE_SIZE = 100; -const MAX_PAGE = 100_000; +const isUsableFilter = (filter: ColumnFiltersState[number]): boolean => + filter.id !== "status" || isKeyStatusFilter(filter.value); -const boundedInteger = (min: number, max: number, fallback: number) => - createParser({ - parse: (value: string) => { - const parsed = parseAsInteger.parse(value); - return parsed === null ? null : Math.min(Math.max(parsed, min), max); - }, - serialize: String, - }).withDefault(fallback); - -// The filters carry a prefix because /api-keys also takes team_id, key_alias and key_type -// as create-key prefills; an unprefixed filter would hijack those deep links. -const TABLE_STATE = { - key_search: parseAsString.withDefault(""), - sort_by: parseAsString.withDefault(DEFAULT_SORT_BY), - sort_order: parseAsStringLiteral(["asc", "desc"] as const).withDefault(DEFAULT_SORT_ORDER), - page: boundedInteger(1, MAX_PAGE, 1), - page_size: boundedInteger(1, MAX_PAGE_SIZE, DEFAULT_PAGE_SIZE), - filter_team: parseAsString.withDefault(""), - filter_org: parseAsString.withDefault(""), - filter_user: parseAsString.withDefault(""), - filter_key_id: parseAsString.withDefault(""), - filter_status: parseAsString.withDefault(""), +const TABLE_STATE_OPTIONS: UrlTableStateOptions = { + sortFields: KEY_TABLE_SORT_FIELDS, + defaultSort: { id: "created_at", desc: true }, + defaultPageSize: 50, + maxPageSize: 100, + filterColumns: FILTER_COLUMNS, + urlKeys: { + search: "key_search", + filter_team_id: "filter_team", + filter_org_id: "filter_org", + filter_user_id: "filter_user", + filter_key_hash: "filter_key_id", + }, }; -const toSortOrder = (active: SortingState[number]): "asc" | "desc" => (active.desc ? "desc" : "asc"); - -const filterValue = (filters: ColumnFiltersState, column: FilterColumn): string | null => { +const appliedFilter = (filters: ColumnFiltersState, column: FilterColumn): string | undefined => { const value = filters.find((filter) => filter.id === column)?.value; - return (typeof value === "string" ? value.trim() : "") || null; + return typeof value === "string" ? value : undefined; }; export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { @@ -103,50 +92,38 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const allTeams = useMemo(() => fetchedTeams ?? [], [fetchedTeams]); const [selectedKeyId, setSelectedKeyId] = useQueryState("key", parseAsString.withOptions({ history: "push" })); - const [tableState, setTableState] = useQueryStates(TABLE_STATE); + const { + search: searchInput, + setSearch, + sorting, + onSortingChange, + pagination, + onPaginationChange, + columnFilters: urlColumnFilters, + onColumnFiltersChange: setUrlColumnFilters, + } = useUrlTableState(TABLE_STATE_OPTIONS); + const columnFilters = useMemo(() => urlColumnFilters.filter(isUsableFilter), [urlColumnFilters]); + const onColumnFiltersChange = useCallback>( + (updaterOrValue) => setUrlColumnFilters(functionalUpdate(updaterOrValue, columnFilters)), + [columnFilters, setUrlColumnFilters], + ); + const { columnVisibility, onColumnVisibilityChange } = usePersistedColumnVisibility( + "virtual-keys", + KEY_TABLE_HIDDEN_COLUMNS, + ); const [filtersOpen, setFiltersOpen] = useState(false); - const searchInput = tableState.key_search; const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); - // A hand-edited sort_by the table cannot sort by would 400 at /key/list and leave the page loading. - const sortBy = KEY_TABLE_SORT_FIELDS.includes(tableState.sort_by) ? tableState.sort_by : DEFAULT_SORT_BY; - const sorting = useMemo( - () => [{ id: sortBy, desc: tableState.sort_order === "desc" }], - [sortBy, tableState.sort_order], - ); - const tablePagination = useMemo( - () => ({ pageIndex: tableState.page - 1, pageSize: tableState.page_size }), - [tableState.page, tableState.page_size], - ); - const { filter_team, filter_org, filter_user, filter_key_id, filter_status } = tableState; - const appliedFilters = useMemo( - () => ({ - team_id: filter_team.trim(), - org_id: filter_org.trim(), - user_id: filter_user.trim(), - key_hash: filter_key_id.trim(), - status: isKeyStatusFilter(filter_status) ? filter_status : "", - }), - [filter_team, filter_org, filter_user, filter_key_id, filter_status], - ); - const columnFilters = useMemo( - () => - FILTER_COLUMNS.filter((column) => appliedFilters[column]).map((column) => ({ - id: column, - value: appliedFilters[column], - })), - [appliedFilters], - ); - + const [activeSort] = sorting; const keyListOptions = { - teamID: appliedFilters.team_id || undefined, - organizationID: appliedFilters.org_id || undefined, + teamID: appliedFilter(columnFilters, "team_id"), + organizationID: appliedFilter(columnFilters, "org_id"), search: searchQuery.trim() || undefined, - userID: appliedFilters.user_id || undefined, - keyHash: appliedFilters.key_hash || undefined, - status: appliedFilters.status || undefined, - sortBy, - sortOrder: tableState.sort_order, + userID: appliedFilter(columnFilters, "user_id"), + keyHash: appliedFilter(columnFilters, "key_hash"), + status: appliedFilter(columnFilters, "status"), + sortBy: activeSort.id, + sortOrder: activeSort.desc ? "desc" : "asc", expand: "user", }; @@ -155,55 +132,13 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { isPending, isPlaceholderData, isFetching, + isError, refetch, - } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, keyListOptions); + } = useKeys(pagination.pageIndex + 1, pagination.pageSize, keyListOptions); const keyList = useMemo(() => keys?.keys ?? [], [keys]); const rowCount = keys?.total_count ?? 0; - const handleSearchChange = useCallback( - (value: string) => { - void setTableState({ key_search: value || null, page: null }); - }, - [setTableState], - ); - - const handleSortingChange = useCallback>( - (updaterOrValue) => { - const active = functionalUpdate(updaterOrValue, sorting)[0]; - void setTableState({ - sort_by: active?.id ?? null, - sort_order: active ? toSortOrder(active) : null, - page: null, - }); - }, - [sorting, setTableState], - ); - - const handleColumnFiltersChange = useCallback>( - (updaterOrValue) => { - const next = functionalUpdate(updaterOrValue, columnFilters); - const nextFilters = { - filter_team: filterValue(next, "team_id"), - filter_org: filterValue(next, "org_id"), - filter_user: filterValue(next, "user_id"), - filter_key_id: filterValue(next, "key_hash"), - filter_status: filterValue(next, "status"), - page: null, - }; - void setTableState(nextFilters); - }, - [columnFilters, setTableState], - ); - - const handlePaginationChange = useCallback>( - (updaterOrValue) => { - const next = functionalUpdate(updaterOrValue, tablePagination); - void setTableState({ page: next.pageIndex + 1, page_size: next.pageSize }); - }, - [tablePagination, setTableState], - ); - const columns = useMemo( () => getKeyTableColumns({ allTeams, organizations, onSelectKey: (key) => void setSelectedKeyId(key.token) }), [allTeams, organizations, setSelectedKeyId], @@ -296,20 +231,22 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { data={keyList} columns={columns} getRowId={(row) => row.token} - defaultColumnVisibility={KEY_TABLE_HIDDEN_COLUMNS} + columnVisibility={columnVisibility} + onColumnVisibilityChange={onColumnVisibilityChange} sortingMode="server" sorting={sorting} - onSortingChange={handleSortingChange} + onSortingChange={onSortingChange} paginationMode="server" - pagination={tablePagination} - onPaginationChange={handlePaginationChange} + pagination={pagination} + onPaginationChange={onPaginationChange} rowCount={rowCount} filterMode="server" columnFilters={columnFilters} - onColumnFiltersChange={handleColumnFiltersChange} + onColumnFiltersChange={onColumnFiltersChange} enableColumnResizing columnResizeMode="onChange" isLoading={isPending || isPlaceholderData} + isError={isError} loadingMessage="Loading keys..." noDataMessage="No keys found" fillHeight @@ -319,7 +256,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { refetch?.()} isRefreshing={isFetching} diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 578e355b85d..3b2a17101ee 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -20,25 +20,39 @@ describe("networking - expired session handling", () => { global.fetch = originalFetch; }); - it("should call clearTokenCookies on expired session", async () => { - const errorData = "Authentication Error - Expired Key"; - const { toast } = await import("@/lib/toast"); + const loadFreshHandleError = async () => { + vi.resetModules(); + const fresh = await import("./networking"); + return fresh.handleError; + }; - if (errorData.includes("Authentication Error - Expired Key")) { - toast.info("UI Session Expired. Logging out."); - clearTokenCookies(); - } + const stubLocation = (pathname: string, search: string, hash: string) => { + const location = { pathname, search, hash, href: "" }; + vi.stubGlobal("window", { location }); + return location; + }; + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("keeps the query string and hash on the redirect after session expiry", async () => { + const handleError = await loadFreshHandleError(); + const location = stubLocation("/ui/api-keys/", "?filter_team=t1&page=2", "#row-3"); + + await handleError("Authentication Error - Expired Key"); + + expect(location.href).toBe("/ui/api-keys/?filter_team=t1&page=2#row-3"); expect(clearTokenCookies).toHaveBeenCalledOnce(); }); - it("should not clear cookies for non-authentication errors", () => { - const errorData = "Some other error"; + it("does not navigate or clear cookies for other errors", async () => { + const handleError = await loadFreshHandleError(); + const location = stubLocation("/ui/api-keys/", "?filter_team=t1&page=2", ""); - if (errorData.includes("Authentication Error - Expired Key")) { - clearTokenCookies(); - } + await handleError("Some other error"); + expect(location.href).toBe(""); expect(clearTokenCookies).not.toHaveBeenCalled(); }); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index cab073dc808..e77c8ba7e41 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -383,7 +383,7 @@ export const handleError = async (errorData: string | any) => { clearTokenCookies(); const browserLocation = getWindowLocation(); if (browserLocation) { - window.location.href = browserLocation.pathname; + window.location.href = browserLocation.pathname + browserLocation.search + browserLocation.hash; } } lastErrorTime = currentTime; diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx index 7bbd4f918cd..c6c6a392aec 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test-d.tsx @@ -1,4 +1,10 @@ -import type { ColumnDef, PaginationState, RowSelectionState, SortingState } from "@tanstack/react-table"; +import type { + ColumnDef, + PaginationState, + RowSelectionState, + SortingState, + VisibilityState, +} from "@tanstack/react-table"; import { DataTable } from "./DataTable"; @@ -12,6 +18,7 @@ const columns: ColumnDef[] = []; const sorting: SortingState = [{ id: "name", desc: false }]; const pagination: PaginationState = { pageIndex: 0, pageSize: 10 }; const rowSelection: RowSelectionState = { r1: true }; +const columnVisibility: VisibilityState = { name: false }; const noop = () => {}; export const uncontrolled = ; @@ -32,6 +39,8 @@ export const controlled = ( onColumnFiltersChange={noop} rowSelection={rowSelection} onRowSelectionChange={noop} + columnVisibility={columnVisibility} + onColumnVisibilityChange={noop} /> ); @@ -65,3 +74,19 @@ export const selectionWithoutHandler = ( // @ts-expect-error a controlled `rowSelection` needs `onRowSelectionChange` or selection changes are dropped ); + +export const visibilityWithoutHandler = ( + // @ts-expect-error a controlled `columnVisibility` needs `onColumnVisibilityChange` or Columns-menu toggles are dropped + +); + +export const bothVisibilitySources = ( + // @ts-expect-error `defaultColumnVisibility` seeds uncontrolled visibility, so it cannot pair with a controlled `columnVisibility` + +); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 8ed8e392ae1..336fc43d695 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -1,4 +1,4 @@ -import type { ColumnDef, ExpandedState, OnChangeFn, PaginationState } from "@tanstack/react-table"; +import type { ColumnDef, ExpandedState, OnChangeFn, PaginationState, VisibilityState } from "@tanstack/react-table"; import { render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { useState } from "react"; @@ -278,11 +278,18 @@ describe("DataTable pagination", () => { type ServerPageHarnessProps = { rowCount: number; isLoading?: boolean; + isError?: boolean; initialPageIndex: number; onChange: (next: PaginationState) => void; }; - function ServerPageHarness({ rowCount, isLoading = false, initialPageIndex, onChange }: ServerPageHarnessProps) { + function ServerPageHarness({ + rowCount, + isLoading = false, + isError = false, + initialPageIndex, + onChange, + }: ServerPageHarnessProps) { const [pagination, setPagination] = useState({ pageIndex: initialPageIndex, pageSize: 10 }); const handleChange: OnChangeFn = (updater) => { const next = typeof updater === "function" ? updater(pagination) : updater; @@ -298,6 +305,7 @@ describe("DataTable pagination", () => { onPaginationChange={handleChange} rowCount={rowCount} isLoading={isLoading} + isError={isError} /> ); } @@ -339,6 +347,29 @@ describe("DataTable pagination", () => { expect(onChange).toHaveBeenCalledTimes(1); expect(screen.getByText("Page 2 of 2")).toBeInTheDocument(); }); + + it("server mode keeps a deep-linked page when the fetch failed, instead of snapping to page 1 on rowCount 0", async () => { + const onChange = vi.fn(); + render(); + + expect(screen.getByText("Page 3 of 1")).toBeInTheDocument(); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(onChange).not.toHaveBeenCalled(); + }); + + it("server mode resumes clamping once the error clears and a real rowCount arrives", async () => { + const onChange = vi.fn(); + const { rerender } = render(); + + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(onChange).not.toHaveBeenCalled(); + + rerender(); + + await waitFor(() => expect(onChange).toHaveBeenCalledWith({ pageIndex: 1, pageSize: 10 })); + expect(onChange).toHaveBeenCalledTimes(1); + expect(screen.getByText("Page 2 of 2")).toBeInTheDocument(); + }); }); describe("DataTable filtering", () => { @@ -555,6 +586,69 @@ describe("DataTable column visibility", () => { expect(await screen.findByTestId("view-option-email")).toBeInTheDocument(); expect(screen.queryByTestId("view-option-name")).not.toBeInTheDocument(); }); + + it("uncontrolled mode seeds hidden columns from defaultColumnVisibility and still toggles internally", async () => { + const user = userEvent.setup(); + render( + } + />, + ); + + expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument(); + await user.click(screen.getByTestId("view-options-trigger")); + await user.click(await screen.findByTestId("view-option-email")); + expect(await screen.findByRole("columnheader", { name: "Email" })).toBeInTheDocument(); + }); + + it("controlled mode hides columns from the prop and reports toggles without changing them locally", async () => { + const user = userEvent.setup(); + const onColumnVisibilityChange = vi.fn>(); + render( + } + />, + ); + + expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument(); + await user.click(screen.getByTestId("view-options-trigger")); + await user.click(await screen.findByTestId("view-option-email")); + + expect(onColumnVisibilityChange).toHaveBeenCalledTimes(1); + const updater = onColumnVisibilityChange.mock.calls[0]?.[0]; + const next = typeof updater === "function" ? updater({ email: false }) : updater; + expect(next).toEqual({ email: true }); + expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument(); + }); + + it("controlled mode reveals the column once the parent applies the reported change", async () => { + const user = userEvent.setup(); + const Harness = () => { + const [columnVisibility, setColumnVisibility] = useState({ email: false }); + return ( + } + /> + ); + }; + render(); + + expect(screen.queryByRole("columnheader", { name: "Email" })).not.toBeInTheDocument(); + await user.click(screen.getByTestId("view-options-trigger")); + await user.click(await screen.findByTestId("view-option-email")); + expect(await screen.findByRole("columnheader", { name: "Email" })).toBeInTheDocument(); + }); }); describe("DataTable pinned columns", () => { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index 26162a3f1f7..e0f57ae1052 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -457,6 +457,7 @@ function useDataTableInstance( onPaginationChange, rowCount, isLoading = false, + isError, pageSizeOptions = DEFAULT_PAGE_SIZE_OPTIONS, filterMode = "none", columnFilters, @@ -466,6 +467,8 @@ function useDataTableInstance( onGlobalFilterChange, enableColumnResizing = false, columnResizeMode = "onEnd", + columnVisibility, + onColumnVisibilityChange, defaultColumnVisibility, getRowCanExpand, renderSubComponent, @@ -481,7 +484,7 @@ function useDataTableInstance( pageIndex: 0, pageSize: pageSizeOptions[0] ?? 25, }); - useServerPageClamp(paginationMode === "server" && !isLoading, rowCount, paginationState); + useServerPageClamp(paginationMode === "server" && !isLoading && !isError, rowCount, paginationState); const filterState = useControllable( columnFilters, onColumnFiltersChange, @@ -490,7 +493,11 @@ function useDataTableInstance( const globalFilterState = useControllable(globalFilter, onGlobalFilterChange, ""); const expandedState = useControllable(expanded, onExpandedChange, {}); const rowSelectionState = useControllable(rowSelection, onRowSelectionChange, {}); - const [columnVisibility, setColumnVisibility] = useState(defaultColumnVisibility ?? {}); + const columnVisibilityState = useControllable( + columnVisibility, + onColumnVisibilityChange, + defaultColumnVisibility ?? {}, + ); const [columnSizing, setColumnSizing] = useState({}); const columnPinning = React.useMemo(() => derivePinning(columns), [columns]); const expansionGuard = renderSubComponent !== undefined ? getRowCanExpand : undefined; @@ -505,7 +512,7 @@ function useDataTableInstance( globalFilter: globalFilterState.value, expanded: expandedState.value, rowSelection: rowSelectionState.value, - columnVisibility, + columnVisibility: columnVisibilityState.value, columnSizing, }, initialState: { columnPinning }, @@ -521,7 +528,7 @@ function useDataTableInstance( onGlobalFilterChange: globalFilterState.onChange, onExpandedChange: expandedState.onChange, onRowSelectionChange: rowSelectionState.onChange, - onColumnVisibilityChange: setColumnVisibility, + onColumnVisibilityChange: columnVisibilityState.onChange, onColumnSizingChange: setColumnSizing, getColumnCanGlobalFilter: (column) => columnCanGlobalFilter(data[0], column), getCoreRowModel: getCoreRowModel(), diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts index 39a887ba948..85cc5f287e0 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/index.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/index.ts @@ -12,6 +12,8 @@ export { type DataTableSortVariant, type DataTableSortField, } from "./DataTableSortHeader"; +export { usePersistedColumnVisibility } from "./usePersistedColumnVisibility"; +export { useUrlTableState, type UrlTableState, type UrlTableStateOptions } from "./useUrlTableState"; export type { DataTablePaginationProps } from "./DataTablePagination"; export type { ColumnPinnedSide, diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts index c767a0a64c0..4529e3df164 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/types.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/types.ts @@ -27,6 +27,7 @@ export interface DataTableResolvedProps { getRowId?: (row: TData, index: number, parent?: Row) => string; isLoading?: boolean; + isError?: boolean; loadingMessage?: string; skeletonRowCount?: number; noDataMessage?: React.ReactNode; @@ -53,6 +54,8 @@ export interface DataTableResolvedProps { enableColumnResizing?: boolean; columnResizeMode?: ColumnResizeMode; + columnVisibility?: VisibilityState; + onColumnVisibilityChange?: OnChangeFn; defaultColumnVisibility?: VisibilityState; getRowCanExpand?: (row: Row) => boolean; @@ -96,6 +99,9 @@ type DataTableBaseProps = Omit< | "columnFilters" | "onColumnFiltersChange" | "defaultColumnFilters" + | "columnVisibility" + | "onColumnVisibilityChange" + | "defaultColumnVisibility" | "rowSelection" | "onRowSelectionChange" >; @@ -142,6 +148,18 @@ type FilterProps = defaultColumnFilters?: ColumnFiltersState; }; +type ColumnVisibilityProps = + | { + columnVisibility: VisibilityState; + onColumnVisibilityChange: OnChangeFn; + defaultColumnVisibility?: never; + } + | { + columnVisibility?: never; + onColumnVisibilityChange?: never; + defaultColumnVisibility?: VisibilityState; + }; + type RowSelectionProps = | { rowSelection: RowSelectionState; onRowSelectionChange: OnChangeFn } | { rowSelection?: never; onRowSelectionChange?: OnChangeFn }; @@ -150,4 +168,5 @@ export type DataTableProps = DataTableBaseProps `litellm_table_columns_${tableId}`; + +const stored = (tableId: string): unknown => { + const raw = localStorage.getItem(keyFor(tableId)); + return raw === null ? null : JSON.parse(raw); +}; + +describe("usePersistedColumnVisibility", () => { + beforeEach(() => { + localStorage.clear(); + }); + + afterEach(() => { + localStorage.clear(); + vi.restoreAllMocks(); + }); + + it("layers the stored choices over the defaults, so a default added after the snapshot still applies", () => { + localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false, spend: true })); + + const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false, name: false })); + + expect(result.current.columnVisibility).toEqual({ email: false, spend: true, name: false }); + }); + + it("falls back to the defaults when nothing is stored, and to {} without defaults", () => { + const withDefaults = renderHook(() => usePersistedColumnVisibility("keys", { spend: false })); + expect(withDefaults.result.current.columnVisibility).toEqual({ spend: false }); + + const bare = renderHook(() => usePersistedColumnVisibility("keys")); + expect(bare.result.current.columnVisibility).toEqual({}); + }); + + it("writes an object update to state and storage", () => { + const { result } = renderHook(() => usePersistedColumnVisibility("keys")); + + act(() => result.current.onColumnVisibilityChange({ email: false })); + + expect(result.current.columnVisibility).toEqual({ email: false }); + expect(stored("keys")).toEqual({ email: false }); + }); + + it("resolves a function updater against the current state before persisting", () => { + localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false })); + const { result } = renderHook(() => usePersistedColumnVisibility("keys")); + + act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, name: false }))); + + expect(result.current.columnVisibility).toEqual({ email: false, name: false }); + expect(stored("keys")).toEqual({ email: false, name: false }); + }); + + it.each([ + ["truncated JSON", '{"email":fal'], + ["a JSON scalar", "42"], + ["a JSON array", "[true]"], + ["non-boolean values", JSON.stringify({ email: "no" })], + ])("falls back to the defaults when storage holds %s", (_label, raw) => { + localStorage.setItem(keyFor("keys"), raw); + + const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false })); + + expect(result.current.columnVisibility).toEqual({ spend: false }); + }); + + it("keeps distinct tableIds isolated in state and storage", () => { + const keys = renderHook(() => usePersistedColumnVisibility("keys")); + const teams = renderHook(() => usePersistedColumnVisibility("teams")); + + act(() => keys.result.current.onColumnVisibilityChange({ email: false })); + + expect(keys.result.current.columnVisibility).toEqual({ email: false }); + expect(teams.result.current.columnVisibility).toEqual({}); + expect(stored("keys")).toEqual({ email: false }); + expect(stored("teams")).toBeNull(); + }); + + it("returns the defaults without throwing when storage is unavailable", () => { + vi.spyOn(console, "warn").mockImplementation(() => {}); + vi.spyOn(Storage.prototype, "getItem").mockImplementation(() => { + throw new Error("SecurityError"); + }); + vi.spyOn(Storage.prototype, "setItem").mockImplementation(() => { + throw new Error("QuotaExceededError"); + }); + + const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false })); + expect(result.current.columnVisibility).toEqual({ spend: false }); + + act(() => result.current.onColumnVisibilityChange({ email: false })); + expect(result.current.columnVisibility).toEqual({ email: false }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts new file mode 100644 index 00000000000..b56ae13d63b --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts @@ -0,0 +1,54 @@ +import type { OnChangeFn, VisibilityState } from "@tanstack/react-table"; +import { useCallback, useState } from "react"; + +import { getLocalStorageItem, setLocalStorageItem } from "@/utils/localStorageUtils"; + +const STORAGE_KEY_PREFIX = "litellm_table_columns_"; + +const EMPTY_VISIBILITY: VisibilityState = {}; + +function storageKey(tableId: string): string { + return `${STORAGE_KEY_PREFIX}${tableId}`; +} + +function isVisibilityState(value: unknown): value is VisibilityState { + if (typeof value !== "object" || value === null || Array.isArray(value)) { + return false; + } + return Object.values(value).every((visible) => typeof visible === "boolean"); +} + +function readStoredVisibility(tableId: string, defaults: VisibilityState): VisibilityState { + const raw = getLocalStorageItem(storageKey(tableId)); + if (raw === null) { + return defaults; + } + try { + const parsed: unknown = JSON.parse(raw); + return isVisibilityState(parsed) ? { ...defaults, ...parsed } : defaults; + } catch { + return defaults; + } +} + +export function usePersistedColumnVisibility( + tableId: string, + defaults: VisibilityState = EMPTY_VISIBILITY, +): { columnVisibility: VisibilityState; onColumnVisibilityChange: OnChangeFn } { + const [columnVisibility, setColumnVisibility] = useState(() => + readStoredVisibility(tableId, defaults), + ); + + const onColumnVisibilityChange = useCallback>( + (updater) => { + setColumnVisibility((previous) => { + const next = typeof updater === "function" ? updater(previous) : updater; + setLocalStorageItem(storageKey(tableId), JSON.stringify(next)); + return next; + }); + }, + [tableId], + ); + + return { columnVisibility, onColumnVisibilityChange }; +} diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx new file mode 100644 index 00000000000..ad46d18b5d8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.test.tsx @@ -0,0 +1,325 @@ +import { SortingState } from "@tanstack/react-table"; +import { act, renderHook, waitFor } from "@testing-library/react"; +import { withNuqsTestingAdapter, type OnUrlUpdateFunction } from "nuqs/adapters/testing"; +import { describe, expect, it, Mock, vi } from "vitest"; +import { useUrlTableState, type UrlTableStateOptions } from "./useUrlTableState"; + +const FILTER_COLUMNS = ["team_id", "user_id"] as const; +type FilterColumn = (typeof FILTER_COLUMNS)[number]; + +const BASE_OPTIONS: UrlTableStateOptions = { + sortFields: ["created_at", "spend", "key_alias"], + defaultSort: { id: "created_at", desc: true }, + defaultPageSize: 50, + filterColumns: FILTER_COLUMNS, +}; + +const PREFIXED_AND_UNPREFIXED_PARAMS = { + audit_page: "2", + audit_page_size: "10", + audit_search: "prefixed", + audit_sort_by: "spend", + audit_sort_order: "asc", + audit_filter_team_id: "team-1", + page: "5", + search: "unprefixed", + filter_team_id: "other-team", +}; + +const RENAMED_AND_DEFAULT_PARAMS = { + key_search: "prod", + filter_team: "team-1", + search: "ignored", + filter_team_id: "ignored", +}; + +const flipDirection = (previous: SortingState): SortingState => previous.map((sort) => ({ ...sort, desc: !sort.desc })); + +const renderTableState = ( + searchParams: Record = {}, + overrides: Partial> = {}, +) => { + const onUrlUpdate = vi.fn(); + const options = { ...BASE_OPTIONS, ...overrides }; + const hook = renderHook(() => useUrlTableState(options), { + wrapper: withNuqsTestingAdapter({ searchParams, onUrlUpdate, hasMemory: true }), + }); + return { ...hook, onUrlUpdate }; +}; + +const lastUrl = (onUrlUpdate: Mock) => { + const event = onUrlUpdate.mock.calls.at(-1)?.[0]; + if (!event) throw new Error("no URL update was emitted"); + return event; +}; + +const flushUrl = async (onUrlUpdate: Mock, write: () => void) => { + const callsBefore = onUrlUpdate.mock.calls.length; + await act(async () => { + write(); + }); + await waitFor(() => expect(onUrlUpdate.mock.calls.length).toBeGreaterThan(callsBefore)); + return lastUrl(onUrlUpdate).searchParams; +}; + +describe("reading table state from the URL", () => { + it("falls back to the defaults when the URL carries no table state", () => { + const { result } = renderTableState(); + + expect(result.current.search).toBe(""); + expect(result.current.sorting).toEqual([{ id: "created_at", desc: true }]); + expect(result.current.pagination).toEqual({ pageIndex: 0, pageSize: 50 }); + expect(result.current.columnFilters).toEqual([]); + }); + + it("maps the 1-based page and page_size onto TanStack pagination", () => { + const { result } = renderTableState({ page: "3", page_size: "25" }); + + expect(result.current.pagination).toEqual({ pageIndex: 2, pageSize: 25 }); + }); + + it.each(["0", "-3", "not-a-number"])("clamps a page of %s up to the first page", (page) => { + const { result } = renderTableState({ page }); + + expect(result.current.pagination.pageIndex).toBe(0); + }); + + it.each([ + ["1000", undefined, 100], + ["1000", 20, 20], + ["0", undefined, 1], + ])("clamps a page_size of %s with maxPageSize %s to %s", (pageSize, maxPageSize, expected) => { + const { result } = renderTableState({ page_size: pageSize }, { maxPageSize }); + + expect(result.current.pagination.pageSize).toBe(expected); + }); + + it("reads a sortable sort_by and its sort_order", () => { + const { result } = renderTableState({ sort_by: "spend", sort_order: "asc" }); + + expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]); + }); + + it("resolves a sort_by outside the allow-list to the default column while keeping the URL's direction", () => { + const { result } = renderTableState({ sort_by: "totally_unknown", sort_order: "asc" }); + + expect(result.current.sorting).toEqual([{ id: "created_at", desc: false }]); + }); + + it("maps filter_ params onto columnFilters, trimming whitespace and dropping blanks", () => { + const { result } = renderTableState({ filter_team_id: "team-1", filter_user_id: " " }); + + expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]); + + const trimmed = renderTableState({ filter_user_id: " user-42 " }); + expect(trimmed.result.current.columnFilters).toEqual([{ id: "user_id", value: "user-42" }]); + }); + + it("reads the search term verbatim so the input can hold trailing spaces", () => { + const { result } = renderTableState({ search: "prod " }); + + expect(result.current.search).toBe("prod "); + }); + + it("reads every key under keyPrefix and ignores the unprefixed ones", () => { + const { result } = renderTableState(PREFIXED_AND_UNPREFIXED_PARAMS, { keyPrefix: "audit_" }); + + expect(result.current.pagination).toEqual({ pageIndex: 1, pageSize: 10 }); + expect(result.current.search).toBe("prefixed"); + expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]); + expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]); + }); + + it("reads renamed keys from urlKeys and ignores the default names", () => { + const { result } = renderTableState(RENAMED_AND_DEFAULT_PARAMS, { + urlKeys: { search: "key_search", filter_team_id: "filter_team" }, + }); + + expect(result.current.search).toBe("prod"); + expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]); + }); + + it("applies keyPrefix in front of a renamed key", () => { + const { result } = renderTableState( + { audit_key_search: "prod", key_search: "ignored" }, + { keyPrefix: "audit_", urlKeys: { search: "key_search" } }, + ); + + expect(result.current.search).toBe("prod"); + }); +}); + +describe("writing table state to the URL", () => { + it("resolves a function updater against the current pagination and replaces history", async () => { + const { result, onUrlUpdate } = renderTableState({ page: "2" }); + + const url = await flushUrl(onUrlUpdate, () => + result.current.onPaginationChange((previous) => ({ ...previous, pageIndex: previous.pageIndex + 1 })), + ); + + expect(url.get("page")).toBe("3"); + expect(url.has("page_size")).toBe(false); + expect(lastUrl(onUrlUpdate).options.history).toBe("replace"); + expect(result.current.pagination).toEqual({ pageIndex: 2, pageSize: 50 }); + }); + + it("writes page_size and drops it again once it returns to the default", async () => { + const { result, onUrlUpdate } = renderTableState(); + + const withSize = await flushUrl(onUrlUpdate, () => + result.current.onPaginationChange({ pageIndex: 0, pageSize: 25 }), + ); + expect(withSize.get("page_size")).toBe("25"); + expect(withSize.has("page")).toBe(false); + + const backToDefault = await flushUrl(onUrlUpdate, () => + result.current.onPaginationChange({ pageIndex: 0, pageSize: 50 }), + ); + expect(backToDefault.has("page_size")).toBe(false); + }); + + it("setSearch writes the term and returns to the first page", async () => { + const { result, onUrlUpdate } = renderTableState({ page: "3" }); + + const url = await flushUrl(onUrlUpdate, () => result.current.setSearch("prod")); + + expect(url.get("search")).toBe("prod"); + expect(url.has("page")).toBe(false); + expect(result.current.search).toBe("prod"); + expect(result.current.pagination.pageIndex).toBe(0); + }); + + it("setSearch with an empty string removes the key", async () => { + const { result, onUrlUpdate } = renderTableState({ search: "prod" }); + + const url = await flushUrl(onUrlUpdate, () => result.current.setSearch("")); + + expect(url.has("search")).toBe(false); + expect(result.current.search).toBe(""); + }); + + it("onSortingChange writes sort_by and sort_order and returns to the first page", async () => { + const { result, onUrlUpdate } = renderTableState({ page: "3" }); + + const url = await flushUrl(onUrlUpdate, () => result.current.onSortingChange([{ id: "spend", desc: false }])); + + expect(url.get("sort_by")).toBe("spend"); + expect(url.get("sort_order")).toBe("asc"); + expect(url.has("page")).toBe(false); + expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]); + }); + + it("onSortingChange drops the keys when the sort matches the default or is cleared", async () => { + const { result, onUrlUpdate } = renderTableState({ sort_by: "spend", sort_order: "asc" }); + + const explicitDefault = await flushUrl(onUrlUpdate, () => + result.current.onSortingChange([{ id: "created_at", desc: true }]), + ); + expect(explicitDefault.has("sort_by")).toBe(false); + expect(explicitDefault.has("sort_order")).toBe(false); + + await flushUrl(onUrlUpdate, () => result.current.onSortingChange([{ id: "key_alias", desc: false }])); + const cleared = await flushUrl(onUrlUpdate, () => result.current.onSortingChange([])); + expect(cleared.has("sort_by")).toBe(false); + expect(cleared.has("sort_order")).toBe(false); + expect(result.current.sorting).toEqual([{ id: "created_at", desc: true }]); + }); + + it("onSortingChange resolves a function updater against the current sort", async () => { + const { result, onUrlUpdate } = renderTableState({ sort_by: "spend" }); + + const url = await flushUrl(onUrlUpdate, () => result.current.onSortingChange(flipDirection)); + + expect(url.get("sort_by")).toBe("spend"); + expect(url.get("sort_order")).toBe("asc"); + expect(result.current.sorting).toEqual([{ id: "spend", desc: false }]); + }); + + it("onColumnFiltersChange writes trimmed filter_ keys and returns to the first page", async () => { + const { result, onUrlUpdate } = renderTableState({ page: "3" }); + + const url = await flushUrl(onUrlUpdate, () => + result.current.onColumnFiltersChange([{ id: "team_id", value: " team-1 " }]), + ); + + expect(url.get("filter_team_id")).toBe("team-1"); + expect(url.has("page")).toBe(false); + expect(result.current.columnFilters).toEqual([{ id: "team_id", value: "team-1" }]); + }); + + it("onColumnFiltersChange removes the key for an empty value and for a filter no longer present", async () => { + const { result, onUrlUpdate } = renderTableState({ filter_team_id: "team-1", filter_user_id: "user-42" }); + + const url = await flushUrl(onUrlUpdate, () => result.current.onColumnFiltersChange([{ id: "team_id", value: "" }])); + + expect(url.has("filter_team_id")).toBe(false); + expect(url.has("filter_user_id")).toBe(false); + expect(result.current.columnFilters).toEqual([]); + }); + + it("onColumnFiltersChange ignores a non-string filter value", async () => { + const { result, onUrlUpdate } = renderTableState({ filter_team_id: "team-1" }); + + const url = await flushUrl(onUrlUpdate, () => + result.current.onColumnFiltersChange([{ id: "team_id", value: ["team-1", "team-2"] }]), + ); + + expect(url.has("filter_team_id")).toBe(false); + }); + + it("onColumnFiltersChange resolves a function updater against the current filters", async () => { + const { result, onUrlUpdate } = renderTableState({ filter_team_id: "team-1" }); + + const url = await flushUrl(onUrlUpdate, () => + result.current.onColumnFiltersChange((previous) => [...previous, { id: "user_id", value: "user-42" }]), + ); + + expect(url.get("filter_team_id")).toBe("team-1"); + expect(url.get("filter_user_id")).toBe("user-42"); + }); + + it("writes prefixed and renamed keys only", async () => { + const { result, onUrlUpdate } = renderTableState( + {}, + { keyPrefix: "audit_", urlKeys: { search: "key_search", filter_team_id: "filter_team" } }, + ); + + await flushUrl(onUrlUpdate, () => result.current.setSearch("prod")); + await flushUrl(onUrlUpdate, () => result.current.onSortingChange([{ id: "spend", desc: false }])); + const url = await flushUrl(onUrlUpdate, () => + result.current.onColumnFiltersChange([{ id: "team_id", value: "team-1" }]), + ); + + expect(url.get("audit_key_search")).toBe("prod"); + expect(url.get("audit_sort_by")).toBe("spend"); + expect(url.get("audit_filter_team")).toBe("team-1"); + expect([...url.keys()].filter((key) => !key.startsWith("audit_"))).toEqual([]); + expect(url.has("audit_search")).toBe(false); + expect(url.has("audit_filter_team_id")).toBe(false); + }); +}); + +describe("referential stability", () => { + it("keeps the TanStack state and the page-clamp handler stable across rerenders while the URL is unchanged", () => { + const { result, rerender } = renderTableState({ page: "2", filter_team_id: "team-1", sort_by: "spend" }); + const first = result.current; + + rerender(); + + expect(result.current.sorting).toBe(first.sorting); + expect(result.current.pagination).toBe(first.pagination); + expect(result.current.columnFilters).toBe(first.columnFilters); + expect(result.current.onPaginationChange).toBe(first.onPaginationChange); + }); + + it("hands out new pagination and untouched sorting after a page change", async () => { + const { result, onUrlUpdate } = renderTableState({ sort_by: "spend" }); + const first = result.current; + + await flushUrl(onUrlUpdate, () => result.current.onPaginationChange({ pageIndex: 4, pageSize: 50 })); + + expect(result.current.pagination).not.toBe(first.pagination); + expect(result.current.pagination.pageIndex).toBe(4); + expect(result.current.sorting).toBe(first.sorting); + }); +}); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts new file mode 100644 index 00000000000..1a423663936 --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/DataTable/useUrlTableState.ts @@ -0,0 +1,232 @@ +import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import { createParser, Nullable, parseAsInteger, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs"; +import { useCallback, useMemo } from "react"; + +const SORT_ORDERS = ["asc", "desc"] as const; +type SortOrder = (typeof SORT_ORDERS)[number]; + +const STANDARD_KEYS = ["search", "sort_by", "sort_order", "page", "page_size"] as const; +type StandardKey = (typeof STANDARD_KEYS)[number]; +type FilterStateKey = `filter_${F}`; +type StateKey = StandardKey | FilterStateKey; + +const MAX_PAGE = 100_000; +const DEFAULT_MAX_PAGE_SIZE = 100; + +export interface UrlTableStateOptions { + sortFields: readonly string[]; + defaultSort: { id: string; desc: boolean }; + defaultPageSize: number; + maxPageSize?: number; + filterColumns: readonly F[]; + keyPrefix?: string; + urlKeys?: Partial, string>>; +} + +export interface UrlTableState { + search: string; + setSearch: (value: string) => void; + sorting: SortingState; + onSortingChange: OnChangeFn; + pagination: PaginationState; + onPaginationChange: OnChangeFn; + columnFilters: ColumnFiltersState; + onColumnFiltersChange: OnChangeFn; +} + +const boundedInteger = (min: number, max: number, fallback: number) => + createParser({ + parse: (value: string) => { + const parsed = parseAsInteger.parse(value); + return parsed === null ? null : Math.min(Math.max(parsed, min), max); + }, + serialize: String, + }).withDefault(fallback); + +const optionalString = parseAsString.withDefault(""); +type OptionalStringParser = typeof optionalString; +const sortOrderParser = (fallback: SortOrder) => parseAsStringLiteral(SORT_ORDERS).withDefault(fallback); + +interface StandardValues { + search: string; + sort_by: string; + sort_order: SortOrder; + page: number; + page_size: number; +} +type FilterValues = Record, string>; +type StandardUpdate = Partial>; +type FilterUpdate = Record, string | null> & Pick, "page">; +type SetTableValues = (update: StandardUpdate | FilterUpdate | null) => Promise; + +interface TableQueryState { + values: StandardValues; + filters: FilterValues; + setValues: SetTableValues; +} + +type TableParsers = { + search: OptionalStringParser; + sort_by: OptionalStringParser; + sort_order: ReturnType; + page: ReturnType; + page_size: ReturnType; +} & Record, OptionalStringParser>; + +const useTableQueryStates = ( + parsers: TableParsers, + urlKeys: Record, string>, +): TableQueryState => { + const [state, setState] = useQueryStates(parsers, { urlKeys }); + return useMemo( + () => ({ + values: state as StandardValues, + filters: state as FilterValues, + setValues: setState as SetTableValues, + }), + [state, setState], + ); +}; + +const filterStateKey = (column: F): FilterStateKey => `filter_${column}`; + +const filterParsers = (filterColumns: readonly F[]) => + Object.fromEntries(filterColumns.map((column) => [filterStateKey(column), optionalString])) as Record< + FilterStateKey, + OptionalStringParser + >; + +const resolveUrlKeys = ( + filterColumns: readonly F[], + keyPrefix: string, + renamed: Partial, string>>, +) => { + const stateKeys: readonly StateKey[] = [ + ...STANDARD_KEYS, + ...filterColumns.map((column) => filterStateKey(column)), + ]; + return Object.fromEntries(stateKeys.map((key) => [key, `${keyPrefix}${renamed[key] ?? key}`])) as Record< + StateKey, + string + >; +}; + +const filterValue = (filters: ColumnFiltersState, column: string): string | null => { + const value = filters.find((filter) => filter.id === column)?.value; + return (typeof value === "string" ? value.trim() : "") || null; +}; + +const filterUpdates = (filterColumns: readonly F[], filters: ColumnFiltersState) => + Object.fromEntries(filterColumns.map((column) => [filterStateKey(column), filterValue(filters, column)])) as Record< + FilterStateKey, + string | null + >; + +const toSortOrder = (active: SortingState[number]): SortOrder => (active.desc ? "desc" : "asc"); + +export function useUrlTableState(options: UrlTableStateOptions): UrlTableState { + const { + sortFields, + defaultSort, + defaultPageSize, + maxPageSize = DEFAULT_MAX_PAGE_SIZE, + filterColumns, + keyPrefix = "", + urlKeys: renamedKeys, + } = options; + const defaultSortId = defaultSort.id; + const defaultSortOrder: SortOrder = defaultSort.desc ? "desc" : "asc"; + + const parsers = useMemo>( + () => ({ + search: optionalString, + sort_by: parseAsString.withDefault(defaultSortId), + sort_order: sortOrderParser(defaultSortOrder), + page: boundedInteger(1, MAX_PAGE, 1), + page_size: boundedInteger(1, maxPageSize, defaultPageSize), + ...filterParsers(filterColumns), + }), + [defaultSortId, defaultSortOrder, defaultPageSize, maxPageSize, filterColumns], + ); + const urlKeys = useMemo( + () => resolveUrlKeys(filterColumns, keyPrefix, renamedKeys ?? {}), + [filterColumns, keyPrefix, renamedKeys], + ); + const { values, filters, setValues } = useTableQueryStates(parsers, urlKeys); + + const sortBy = sortFields.includes(values.sort_by) ? values.sort_by : defaultSortId; + const sortDesc = values.sort_order === "desc"; + const sorting = useMemo(() => [{ id: sortBy, desc: sortDesc }], [sortBy, sortDesc]); + + const pagination = useMemo( + () => ({ pageIndex: values.page - 1, pageSize: values.page_size }), + [values.page, values.page_size], + ); + + const columnFilters = useMemo( + () => + filterColumns.flatMap((column) => { + const value = filters[filterStateKey(column)].trim(); + return value ? [{ id: column, value }] : []; + }), + [filterColumns, filters], + ); + + const setSearch = useCallback( + (value: string) => { + void setValues({ search: value || null, page: null }); + }, + [setValues], + ); + + const onSortingChange = useCallback>( + (updaterOrValue) => { + const active = functionalUpdate(updaterOrValue, sorting)[0]; + void setValues({ + sort_by: active?.id ?? null, + sort_order: active ? toSortOrder(active) : null, + page: null, + }); + }, + [setValues, sorting], + ); + + const onPaginationChange = useCallback>( + (updaterOrValue) => { + const next = functionalUpdate(updaterOrValue, pagination); + void setValues({ page: next.pageIndex + 1, page_size: next.pageSize }); + }, + [pagination, setValues], + ); + + const onColumnFiltersChange = useCallback>( + (updaterOrValue) => { + const next = functionalUpdate(updaterOrValue, columnFilters); + void setValues({ ...filterUpdates(filterColumns, next), page: null }); + }, + [columnFilters, filterColumns, setValues], + ); + + return useMemo( + () => ({ + search: values.search, + setSearch, + sorting, + onSortingChange, + pagination, + onPaginationChange, + columnFilters, + onColumnFiltersChange, + }), + [ + values.search, + setSearch, + sorting, + onSortingChange, + pagination, + onPaginationChange, + columnFilters, + onColumnFiltersChange, + ], + ); +} diff --git a/ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx b/ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx new file mode 100644 index 00000000000..427d454ee2d --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useUrlTab.test.tsx @@ -0,0 +1,100 @@ +import { act, renderHook, waitFor } from "@testing-library/react"; +import { NuqsTestingAdapter, type OnUrlUpdateFunction } from "nuqs/adapters/testing"; +import type { ReactNode } from "react"; +import { describe, expect, it, vi } from "vitest"; +import { useUrlTab } from "./useUrlTab"; + +const TABS = ["chat", "compare", "compliance"] as const; +type Tab = (typeof TABS)[number]; + +interface RenderArgs { + searchParams?: string; + onUrlUpdate?: OnUrlUpdateFunction; + key?: string; +} + +const initialProps: { values: readonly Tab[] } = { values: TABS }; + +const renderUrlTab = ({ searchParams, onUrlUpdate, key }: RenderArgs = {}) => + renderHook(({ values }: { values: readonly Tab[] }) => useUrlTab(values, "chat", key), { + initialProps, + wrapper: ({ children }: { children: ReactNode }) => ( + + {children} + + ), + }); + +const lastUrlUpdate = (onUrlUpdate: ReturnType>) => + onUrlUpdate.mock.calls.at(-1)?.[0]; + +describe("useUrlTab", () => { + it("reads the active tab from the URL", () => { + const { result } = renderUrlTab({ searchParams: "?tab=compare" }); + + expect(result.current[0]).toBe("compare"); + }); + + it("resolves a URL value outside the allowed tabs to the fallback and drops it from the URL", async () => { + const onUrlUpdate = vi.fn(); + const { result } = renderUrlTab({ searchParams: "?tab=settings&other=1", onUrlUpdate }); + + expect(result.current[0]).toBe("chat"); + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + expect(lastUrlUpdate(onUrlUpdate)?.searchParams.has("tab")).toBe(false); + expect(lastUrlUpdate(onUrlUpdate)?.searchParams.get("other")).toBe("1"); + }); + + it("leaves a URL that names an allowed tab untouched", async () => { + const onUrlUpdate = vi.fn(); + renderUrlTab({ searchParams: "?tab=compare", onUrlUpdate }); + + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(onUrlUpdate).not.toHaveBeenCalled(); + }); + + it("reads from the caller's key instead of the default one", () => { + const { result } = renderUrlTab({ searchParams: "?view=compliance&tab=compare", key: "view" }); + + expect(result.current[0]).toBe("compliance"); + }); + + it("writes ?tab= with history replace when a tab is selected", async () => { + const onUrlUpdate = vi.fn(); + const { result } = renderUrlTab({ onUrlUpdate }); + + act(() => result.current[1]("compare")); + + await waitFor(() => expect(lastUrlUpdate(onUrlUpdate)?.searchParams.get("tab")).toBe("compare")); + expect(lastUrlUpdate(onUrlUpdate)?.options.history).toBe("replace"); + expect(result.current[0]).toBe("compare"); + }); + + it("removes the param when the fallback tab is selected", async () => { + const onUrlUpdate = vi.fn(); + const { result } = renderUrlTab({ searchParams: "?tab=compare", onUrlUpdate }); + + act(() => result.current[1]("chat")); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + expect(lastUrlUpdate(onUrlUpdate)?.searchParams.has("tab")).toBe(false); + expect(result.current[0]).toBe("chat"); + }); + + it("falls back and clears the param when the current tab is no longer among the allowed values", async () => { + const onUrlUpdate = vi.fn(); + const { result, rerender } = renderUrlTab({ searchParams: "?tab=compliance", onUrlUpdate }); + expect(result.current[0]).toBe("compliance"); + + rerender({ values: ["chat", "compare"] }); + + expect(result.current[0]).toBe("chat"); + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + expect(lastUrlUpdate(onUrlUpdate)?.searchParams.has("tab")).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/hooks/useUrlTab.ts b/ui/litellm-dashboard/src/hooks/useUrlTab.ts new file mode 100644 index 00000000000..2f3d705c610 --- /dev/null +++ b/ui/litellm-dashboard/src/hooks/useUrlTab.ts @@ -0,0 +1,12 @@ +import { parseAsString, useQueryState } from "nuqs"; +import { useCallback, useEffect } from "react"; + +export function useUrlTab(values: readonly T[], fallback: T, key = "tab"): [T, (tab: T) => void] { + const [urlTab, setUrlTab] = useQueryState(key, parseAsString.withDefault(fallback)); + const tab = values.find((value) => value === urlTab) ?? fallback; + useEffect(() => { + if (urlTab !== tab) void setUrlTab(null); + }, [urlTab, tab, setUrlTab]); + const setTab = useCallback((next: T) => void setUrlTab(next), [setUrlTab]); + return [tab, setTab]; +} diff --git a/ui/litellm-dashboard/src/utils/tabRoutes.test.ts b/ui/litellm-dashboard/src/utils/tabRoutes.test.ts deleted file mode 100644 index 402be55c33a..00000000000 --- a/ui/litellm-dashboard/src/utils/tabRoutes.test.ts +++ /dev/null @@ -1,47 +0,0 @@ -/* @vitest-environment jsdom */ -import { describe, expect, it, vi } from "vitest"; - -vi.mock("@/components/networking", () => ({ serverRootPath: "" })); - -import { createTabRoutes } from "./tabRoutes"; - -const routes = createTabRoutes("logs", ["audit", "deleted-keys", "deleted-teams"] as const); - -describe("createTabRoutes.slugFromPathname", () => { - it("returns empty string for the base path with or without a trailing slash", () => { - expect(routes.slugFromPathname("/logs")).toBe(""); - expect(routes.slugFromPathname("/logs/")).toBe(""); - }); - - it("extracts the tab slug from dev and proxy-mounted (/ui) paths", () => { - expect(routes.slugFromPathname("/logs/audit")).toBe("audit"); - expect(routes.slugFromPathname("/ui/logs/deleted-teams/")).toBe("deleted-teams"); - }); - - it("returns the raw segment for an unknown tab so the caller can redirect to base", () => { - expect(routes.slugFromPathname("/ui/logs/bogus")).toBe("bogus"); - }); - - it("returns empty string when the base segment is not in the path", () => { - expect(routes.slugFromPathname("/teams")).toBe(""); - }); -}); - -describe("createTabRoutes.tabHref", () => { - it("builds the trailing-slash base href for the empty slug", () => { - expect(routes.tabHref("")).toBe("/ui/logs/"); - }); - - it("builds a trailing-slash href for every tab slug (required by static export)", () => { - for (const slug of routes.slugs) { - expect(routes.tabHref(slug)).toBe(`/ui/logs/${slug}/`); - } - }); -}); - -describe("createTabRoutes metadata", () => { - it("preserves the base segment and slug tuple", () => { - expect(routes.baseSegment).toBe("logs"); - expect(routes.slugs).toEqual(["audit", "deleted-keys", "deleted-teams"]); - }); -}); diff --git a/ui/litellm-dashboard/src/utils/tabRoutes.ts b/ui/litellm-dashboard/src/utils/tabRoutes.ts deleted file mode 100644 index 4af2b983cba..00000000000 --- a/ui/litellm-dashboard/src/utils/tabRoutes.ts +++ /dev/null @@ -1,26 +0,0 @@ -import { uiHref } from "@/utils/uiHref"; - -export interface TabRoutes { - baseSegment: string; - slugs: readonly Slug[]; - tabHref: (slug: string) => string; - slugFromPathname: (pathname: string) => string; -} - -export function createTabRoutes(baseSegment: string, slugs: readonly Slug[]): TabRoutes { - const tabHref = (slug: string): string => { - const base = uiHref(baseSegment); - return slug ? `${base}/${slug}/` : `${base}/`; - }; - - const slugFromPathname = (pathname: string): string => { - const parts = pathname.split("/").filter(Boolean); - const idx = parts.indexOf(baseSegment); - if (idx === -1) { - return ""; - } - return parts[idx + 1] ?? ""; - }; - - return { baseSegment, slugs, tabHref, slugFromPathname }; -} From 4d30bbce453a06e850cbf6007a42dd7e00803380 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 09:34:05 -0700 Subject: [PATCH 156/187] fix(ui): keep controlled client-side table pages across data reloads Controlled client-mode DataTables no longer let TanStack reset the page index when rows change, since the owner of the pagination state decides the page. Once rows settle, a page past the end snaps back to the last page, matching server mode --- .../shared/DataTable/DataTable.test.tsx | 73 +++++++++++++++++++ .../components/shared/DataTable/DataTable.tsx | 34 ++++++++- 2 files changed, 104 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 336fc43d695..a7e4befa6ac 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -357,6 +357,79 @@ describe("DataTable pagination", () => { expect(onChange).not.toHaveBeenCalled(); }); + type ClientPageHarnessProps = { + data: Person[]; + isLoading?: boolean; + initialPageIndex: number; + onChange: (next: PaginationState) => void; + }; + + function ClientPageHarness({ data, isLoading = false, initialPageIndex, onChange }: ClientPageHarnessProps) { + const [pagination, setPagination] = useState({ pageIndex: initialPageIndex, pageSize: 2 }); + const handleChange: OnChangeFn = (updater) => { + const next = typeof updater === "function" ? updater(pagination) : updater; + onChange(next); + setPagination(next); + }; + return ( + + ); + } + + it("client mode keeps a controlled page when rows arrive after loading and when they are refetched", async () => { + const onChange = vi.fn(); + const { rerender } = render(); + + rerender(); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(names()).toEqual(["P2", "P3"]); + + rerender(); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(names()).toEqual(["P2", "P3"]); + expect(onChange).not.toHaveBeenCalled(); + }); + + it("client mode snaps a controlled page past the end back to the last page", async () => { + const onChange = vi.fn(); + render(); + + await waitFor(() => expect(onChange).toHaveBeenCalledWith({ pageIndex: 2, pageSize: 2 })); + expect(onChange).toHaveBeenCalledTimes(1); + expect(names()).toEqual(["P4"]); + }); + + it("client mode leaves a controlled page alone while there are no rows to page through", async () => { + const onChange = vi.fn(); + render(); + + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(onChange).not.toHaveBeenCalled(); + }); + + it("client mode without a controlled page still returns to the first page when the rows change", async () => { + const user = userEvent.setup(); + const { rerender } = render( + , + ); + + await user.click(screen.getByTestId("pagination-next")); + expect(names()).toEqual(["P2", "P3"]); + + rerender( + , + ); + await waitFor(() => expect(names()).toEqual(["P0", "P1"])); + }); + it("server mode resumes clamping once the error clears and a real rowCount arrives", async () => { const onChange = vi.fn(); const { rerender } = render(); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index e0f57ae1052..340f8d4f44f 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -425,7 +425,7 @@ function useControllable( return { value: internal, onChange: setInternal }; } -function useServerPageClamp( +function usePageClamp( active: boolean, rowCount: number | undefined, pagination: { value: PaginationState; onChange: OnChangeFn }, @@ -484,7 +484,6 @@ function useDataTableInstance( pageIndex: 0, pageSize: pageSizeOptions[0] ?? 25, }); - useServerPageClamp(paginationMode === "server" && !isLoading && !isError, rowCount, paginationState); const filterState = useControllable( columnFilters, onColumnFiltersChange, @@ -536,9 +535,38 @@ function useDataTableInstance( ...(getRowId !== undefined ? { getRowId } : {}), ...(enableRowSelection !== undefined ? { enableRowSelection } : {}), ...(paginationMode === "server" && rowCount !== undefined ? { rowCount } : {}), + autoResetPageIndex: pagination === undefined && paginationMode !== "server", }; - return useReactTable(tableOptions); + const table = useReactTable(tableOptions); + const clampOptions: SettledPageClampOptions = { + paginationMode, + controlled: pagination !== undefined, + settled: !isLoading && !isError, + rowCount, + pagination: paginationState, + }; + useSettledPageClamp(table, clampOptions); + return table; +} + +type SettledPageClampOptions = { + paginationMode: PaginationMode; + controlled: boolean; + settled: boolean; + rowCount: number | undefined; + pagination: { value: PaginationState; onChange: OnChangeFn }; +}; + +function useSettledPageClamp(table: Table, options: SettledPageClampOptions): void { + const { paginationMode, controlled, settled, rowCount, pagination } = options; + const clientRowCount = paginationMode === "client" ? table.getPrePaginationRowModel().rows.length : 0; + const clientPageIsClampable = paginationMode === "client" && controlled && clientRowCount > 0; + usePageClamp( + settled && (paginationMode === "server" || clientPageIsClampable), + paginationMode === "server" ? rowCount : clientRowCount, + pagination, + ); } export function DataTable(props: DataTableProps) { From adc937c4931021c68864ed28014d51b851ecaade Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 12:50:08 -0700 Subject: [PATCH 157/187] fix(ui): read persisted column visibility from storage instead of a mounted copy usePersistedColumnVisibility kept a useState copy seeded once at mount, so a later tableId or defaults change showed the old table's columns and saved them under the new key. It now reads localStorage through useSyncExternalStore, keeping only writes that storage refused in memory, so the hook has no copy to go stale. Stored choices are layered over the defaults on every read, and changes saved in another tab show up. --- .../usePersistedColumnVisibility.test.tsx | 74 ++++++++++++++++++- .../DataTable/usePersistedColumnVisibility.ts | 54 +++++++++++--- 2 files changed, 112 insertions(+), 16 deletions(-) diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx index 8d2e0f61f54..fbfb71e5bc2 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx @@ -1,3 +1,4 @@ +import type { VisibilityState } from "@tanstack/react-table"; import { act, renderHook } from "@testing-library/react"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; @@ -10,6 +11,9 @@ const stored = (tableId: string): unknown => { return raw === null ? null : JSON.parse(raw); }; +const showEveryColumn = (previous: VisibilityState): VisibilityState => + Object.fromEntries(Object.keys(previous).map((column) => [column, true])); + describe("usePersistedColumnVisibility", () => { beforeEach(() => { localStorage.clear(); @@ -55,6 +59,15 @@ describe("usePersistedColumnVisibility", () => { expect(stored("keys")).toEqual({ email: false, name: false }); }); + it("hands a function updater the default-hidden columns, so showing every column sticks", () => { + const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false })); + + act(() => result.current.onColumnVisibilityChange(showEveryColumn)); + + expect(result.current.columnVisibility).toEqual({ spend: true }); + expect(stored("keys")).toEqual({ spend: true }); + }); + it.each([ ["truncated JSON", '{"email":fal'], ["a JSON scalar", "42"], @@ -80,19 +93,72 @@ describe("usePersistedColumnVisibility", () => { expect(stored("teams")).toBeNull(); }); + it("reads and writes the new table's columns after the tableId changes", () => { + localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false })); + localStorage.setItem(keyFor("teams"), JSON.stringify({ spend: false })); + const { result, rerender } = renderHook(({ tableId }) => usePersistedColumnVisibility(tableId), { + initialProps: { tableId: "keys" }, + }); + + rerender({ tableId: "teams" }); + expect(result.current.columnVisibility).toEqual({ spend: false }); + + act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, name: false }))); + expect(stored("teams")).toEqual({ spend: false, name: false }); + expect(stored("keys")).toEqual({ email: false }); + }); + + it("applies new defaults passed after mount", () => { + const { result, rerender } = renderHook(({ defaults }) => usePersistedColumnVisibility("keys", defaults), { + initialProps: { defaults: { spend: false } }, + }); + + rerender({ defaults: { name: false } }); + + expect(result.current.columnVisibility).toEqual({ name: false }); + }); + + it("shows a change another tab saved for the same table", () => { + const { result } = renderHook(() => usePersistedColumnVisibility("keys")); + + act(() => { + localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false })); + window.dispatchEvent(new StorageEvent("storage", { key: keyFor("keys") })); + }); + + expect(result.current.columnVisibility).toEqual({ email: false }); + }); + + it("keeps a toggle that storage refused, and saves the next one once storage accepts it", () => { + localStorage.setItem(keyFor("full"), JSON.stringify({ spend: false })); + vi.spyOn(console, "warn").mockImplementation(() => {}); + vi.spyOn(Storage.prototype, "setItem").mockImplementationOnce(() => { + throw new Error("QuotaExceededError"); + }); + const { result } = renderHook(() => usePersistedColumnVisibility("full")); + + act(() => result.current.onColumnVisibilityChange({ email: false })); + expect(result.current.columnVisibility).toEqual({ email: false }); + expect(stored("full")).toEqual({ spend: false }); + + act(() => result.current.onColumnVisibilityChange({ name: false })); + expect(result.current.columnVisibility).toEqual({ name: false }); + expect(stored("full")).toEqual({ name: false }); + }); + it("returns the defaults without throwing when storage is unavailable", () => { vi.spyOn(console, "warn").mockImplementation(() => {}); vi.spyOn(Storage.prototype, "getItem").mockImplementation(() => { throw new Error("SecurityError"); }); vi.spyOn(Storage.prototype, "setItem").mockImplementation(() => { - throw new Error("QuotaExceededError"); + throw new Error("SecurityError"); }); - const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false })); + const { result } = renderHook(() => usePersistedColumnVisibility("blocked", { spend: false })); expect(result.current.columnVisibility).toEqual({ spend: false }); - act(() => result.current.onColumnVisibilityChange({ email: false })); - expect(result.current.columnVisibility).toEqual({ email: false }); + act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, email: false }))); + expect(result.current.columnVisibility).toEqual({ spend: false, email: false }); }); }); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts index b56ae13d63b..7fffc50b381 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts @@ -1,16 +1,46 @@ import type { OnChangeFn, VisibilityState } from "@tanstack/react-table"; -import { useCallback, useState } from "react"; +import { useCallback, useMemo, useSyncExternalStore } from "react"; -import { getLocalStorageItem, setLocalStorageItem } from "@/utils/localStorageUtils"; +import { + LOCAL_STORAGE_EVENT, + emitLocalStorageChange, + getLocalStorageItem, + setLocalStorageItem, +} from "@/utils/localStorageUtils"; const STORAGE_KEY_PREFIX = "litellm_table_columns_"; const EMPTY_VISIBILITY: VisibilityState = {}; +const unsavedWrites = new Map(); + function storageKey(tableId: string): string { return `${STORAGE_KEY_PREFIX}${tableId}`; } +function subscribe(onChange: () => void): () => void { + window.addEventListener("storage", onChange); + window.addEventListener(LOCAL_STORAGE_EVENT, onChange); + return () => { + window.removeEventListener("storage", onChange); + window.removeEventListener(LOCAL_STORAGE_EVENT, onChange); + }; +} + +function readRaw(key: string): string | null { + return unsavedWrites.get(key) ?? getLocalStorageItem(key); +} + +function writeRaw(key: string, raw: string): void { + setLocalStorageItem(key, raw); + if (getLocalStorageItem(key) === raw) { + unsavedWrites.delete(key); + } else { + unsavedWrites.set(key, raw); + } + emitLocalStorageChange(key); +} + function isVisibilityState(value: unknown): value is VisibilityState { if (typeof value !== "object" || value === null || Array.isArray(value)) { return false; @@ -18,8 +48,7 @@ function isVisibilityState(value: unknown): value is VisibilityState { return Object.values(value).every((visible) => typeof visible === "boolean"); } -function readStoredVisibility(tableId: string, defaults: VisibilityState): VisibilityState { - const raw = getLocalStorageItem(storageKey(tableId)); +function parseVisibility(raw: string | null, defaults: VisibilityState): VisibilityState { if (raw === null) { return defaults; } @@ -35,19 +64,20 @@ export function usePersistedColumnVisibility( tableId: string, defaults: VisibilityState = EMPTY_VISIBILITY, ): { columnVisibility: VisibilityState; onColumnVisibilityChange: OnChangeFn } { - const [columnVisibility, setColumnVisibility] = useState(() => - readStoredVisibility(tableId, defaults), + const key = storageKey(tableId); + const raw = useSyncExternalStore( + subscribe, + () => readRaw(key), + () => null, ); + const columnVisibility = useMemo(() => parseVisibility(raw, defaults), [raw, defaults]); const onColumnVisibilityChange = useCallback>( (updater) => { - setColumnVisibility((previous) => { - const next = typeof updater === "function" ? updater(previous) : updater; - setLocalStorageItem(storageKey(tableId), JSON.stringify(next)); - return next; - }); + const next = typeof updater === "function" ? updater(parseVisibility(readRaw(key), defaults)) : updater; + writeRaw(key, JSON.stringify(next)); }, - [tableId], + [key, defaults], ); return { columnVisibility, onColumnVisibilityChange }; From d29753c52172b619a2ad7702b34fe1e09b45a6f5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 15:43:21 -0700 Subject: [PATCH 158/187] fix(ui): let another tab's column save replace a toggle this tab could not save A column toggle that localStorage refused was kept in memory and read ahead of storage, so a later save from another tab stayed hidden until this tab saved again. A storage event now drops the in-memory copy for its key, or all of them when another tab clears storage --- .../usePersistedColumnVisibility.test.tsx | 32 ++++++++++++++++++- .../DataTable/usePersistedColumnVisibility.ts | 16 ++++++++-- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx index fbfb71e5bc2..35bba906afc 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.test.tsx @@ -109,8 +109,9 @@ describe("usePersistedColumnVisibility", () => { }); it("applies new defaults passed after mount", () => { + const initialProps: { defaults: VisibilityState } = { defaults: { spend: false } }; const { result, rerender } = renderHook(({ defaults }) => usePersistedColumnVisibility("keys", defaults), { - initialProps: { defaults: { spend: false } }, + initialProps, }); rerender({ defaults: { name: false } }); @@ -146,6 +147,35 @@ describe("usePersistedColumnVisibility", () => { expect(stored("full")).toEqual({ name: false }); }); + it("shows another tab's save over a toggle this tab could not save", () => { + vi.spyOn(console, "warn").mockImplementation(() => {}); + vi.spyOn(Storage.prototype, "setItem").mockImplementationOnce(() => { + throw new Error("QuotaExceededError"); + }); + const { result } = renderHook(() => usePersistedColumnVisibility("shadowed")); + act(() => result.current.onColumnVisibilityChange({ email: false })); + + act(() => { + localStorage.setItem(keyFor("shadowed"), JSON.stringify({ name: false })); + window.dispatchEvent(new StorageEvent("storage", { key: keyFor("shadowed") })); + }); + + expect(result.current.columnVisibility).toEqual({ name: false }); + }); + + it("drops a toggle this tab could not save once another tab clears storage", () => { + vi.spyOn(console, "warn").mockImplementation(() => {}); + vi.spyOn(Storage.prototype, "setItem").mockImplementationOnce(() => { + throw new Error("QuotaExceededError"); + }); + const { result } = renderHook(() => usePersistedColumnVisibility("cleared", { spend: false })); + act(() => result.current.onColumnVisibilityChange({ email: false })); + + act(() => window.dispatchEvent(new StorageEvent("storage", { key: null }))); + + expect(result.current.columnVisibility).toEqual({ spend: false }); + }); + it("returns the defaults without throwing when storage is unavailable", () => { vi.spyOn(console, "warn").mockImplementation(() => {}); vi.spyOn(Storage.prototype, "getItem").mockImplementation(() => { diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts index 7fffc50b381..3cbf5c2a000 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts +++ b/ui/litellm-dashboard/src/components/shared/DataTable/usePersistedColumnVisibility.ts @@ -18,11 +18,23 @@ function storageKey(tableId: string): string { return `${STORAGE_KEY_PREFIX}${tableId}`; } +function forgetUnsavedWrite(event: StorageEvent): void { + if (event.key === null) { + unsavedWrites.clear(); + return; + } + unsavedWrites.delete(event.key); +} + function subscribe(onChange: () => void): () => void { - window.addEventListener("storage", onChange); + const onStorage = (event: StorageEvent): void => { + forgetUnsavedWrite(event); + onChange(); + }; + window.addEventListener("storage", onStorage); window.addEventListener(LOCAL_STORAGE_EVENT, onChange); return () => { - window.removeEventListener("storage", onChange); + window.removeEventListener("storage", onStorage); window.removeEventListener(LOCAL_STORAGE_EVENT, onChange); }; } From 36eb9cdb356e0ad02124e7ae149324def11b4669 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 15:56:27 -0700 Subject: [PATCH 159/187] test(proxy): drop the route-list membership test that the PATCH gate tests already cover --- .../proxy/auth/test_route_checks.py | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 3211e85ff97..72c7011e6f5 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -2892,7 +2892,7 @@ def test_team_update_gate_allows_org_admin_with_resolved_org(): ) -def test_team_update_gate_admits_internal_user_without_org_context(): +def test_team_update_gate_admits_internal_user_without_org_context(): # test-quality-ok: the gate's only success signal is not raising; the handler's team-admin 403s are pinned in test_team_endpoints """/team/update is self-managed (LIT-5722): the coarse gate admits any authenticated caller and update_team resolves proxy, org or team admin itself, then filters team admins through the team_admin_editable_team_fields setting. Before that the gate 401'd every @@ -2997,23 +2997,6 @@ async def test_add_team_org_context_noop_for_static_team_route(): assert out == body -def test_patch_team_route_has_same_reach_as_team_update(): - """/team/{team_id} is reachable by org admins (in org_admin_allowed_routes) but - NOT by regular internal users or the role-agnostic self_managed_routes — the - latter would open /team/new (the collision footgun) to any authenticated user.""" - from litellm.proxy._types import LiteLLMRoutes - - assert RouteChecks.check_route_access( - route="/team/abc-123", allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value - ) - assert not RouteChecks.check_route_access( - route="/team/abc-123", allowed_routes=LiteLLMRoutes.internal_user_routes.value - ) - assert not RouteChecks.check_route_access( - route="/team/abc-123", allowed_routes=LiteLLMRoutes.self_managed_routes.value - ) - - def _patch_team_request() -> MagicMock: request = MagicMock(spec=Request) request.method = "PATCH" From afb28540bbd3868dcebd83e8e7c7347c7611abfa Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 16:01:32 -0700 Subject: [PATCH 160/187] fix(e2e): keep the CLI determinism test out of the in-cluster suite It drives the real CLI for several seconds. The edge stamps every upstream call with PYTEST_CURRENT_TEST, a process-global that names whichever test the worker is in when the call arrives rather than the one that made it, so a test that holds a worker that long collects other tests' in-flight calls. Build 234's key report credits this test with 20 Bedrock and 7 Anthropic misses, and it makes no provider call at all. Those misattributed calls take the wrong test id into the cache key and write recordings under it, so the test was polluting the shared corpus it exists to protect. Deselected unless E2E_CLI_DETERMINISM is set, the same opt-in shape the managed-files, prompt-caching and redis-chaos markers already use. The attribution bug itself is older than this branch and is reported, not fixed here. --- .../_driver_unit_tests/test_request_determinism.py | 2 ++ tests/e2e/conftest.py | 7 +++++++ tests/e2e/e2e_config.py | 1 + tests/e2e/pytest.ini | 1 + 4 files changed, 11 insertions(+) diff --git a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py index 5046f35c73b..b7d330b7da6 100644 --- a/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py +++ b/tests/e2e/claude_code/_driver_unit_tests/test_request_determinism.py @@ -36,6 +36,8 @@ import pytest from claude_code.cli_driver import _FIXED_CLI_USER_ID, _seed_cli_identity, _stable_cli_state, run_claude from claude_code.rate_limiter import RateLimiter +pytestmark = pytest.mark.cli_determinism + _STUB_REPLY = { "id": "msg_stub", "type": "message", diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 430e16525d5..ac4cfb71407 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -23,6 +23,7 @@ from typing import Final import pytest import requests from e2e_config import ( + CLI_DETERMINISM_OPT_IN_ENV, CONTROL_PLANE_BASE_URL, FIXTURE_DIR, FIXTURE_MODE_RAW, @@ -53,6 +54,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "managed_files": MANAGED_FILES_OPT_IN_ENV, "prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV, "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, + "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, } ) @@ -120,6 +122,11 @@ def pytest_configure(config: pytest.Config) -> None: "prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including " "prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set", ) + config.addinivalue_line( + "markers", + "cli_determinism: drives the real claude CLI for several seconds, which widens the window in which " + "another test's in-flight upstream call is attributed to it; deselected unless E2E_CLI_DETERMINISM is set", + ) config.addinivalue_line( "markers", "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 896cb3e7efe..82ddb09f7f5 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -145,6 +145,7 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" +CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 1fdd3bd28ad..f6d23a3ec12 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -10,4 +10,5 @@ markers = weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set + cli_determinism: drives the real claude CLI for several seconds, which widens the window in which another test's in-flight upstream call is attributed to it; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set From 02d9aae8c87fd2948a37c47815995a16465e45ce Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:01:52 -0700 Subject: [PATCH 161/187] fix(proxy): stop forwarding LiteLLM credential headers on Bedrock agent-runtime passthrough The agent-runtime branch of /bedrock/{endpoint} (agents, knowledgebases, flows, retrieveAndGenerate, rerank, generateQuery, optimize-prompt) forwarded every caller header to AWS next to the SigV4 signature, so a LiteLLM key presented in x-api-key or x-litellm-api-key reached bedrock-agent-runtime verbatim. Build the upstream header set explicitly: drop LiteLLM credential headers by name and any authenticated secret by value, keep the remaining caller headers, and let the signed headers win on collisions. --- .../llm_passthrough_endpoints.py | 17 ++- .../test_llm_pass_through_endpoints.py | 140 ++++++++++++++++++ 2 files changed, 155 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 0fe9d1cc626..b9b8cb3a22b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1180,9 +1180,8 @@ async def bedrock_proxy_route( endpoint_func: Final = create_pass_through_route( endpoint=endpoint, target=str(prepped.url), - custom_headers=prepped.headers, + custom_headers=_upstream_headers_for_bedrock_agent_runtime_route(request, user_api_key_dict, prepped.headers), is_streaming_request=is_streaming_request, - _forward_headers=True, ) # dynamically construct pass-through endpoint based on incoming path setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, data) # SigV4 signs an exact payload; pass-through must send prepped.body, not json.dumps @@ -2001,6 +2000,9 @@ _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS: Final = frozenset({"authorization", "x-a _HEADERS_NEVER_FORWARDED_TO_ANTHROPIC: Final = frozenset({"content-length", "host", "accept-encoding"}) | ( SpecialHeaders.litellm_credential_header_names() - _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS ) +_HEADERS_NEVER_FORWARDED_TO_BEDROCK: Final = ( + frozenset({"content-length", "host", "accept-encoding"}) | SpecialHeaders.litellm_credential_header_names() +) _MAPPED_ROUTE_CALLER_KEY_HEADER: Final = "litellm_user_api_key" @@ -2099,6 +2101,17 @@ def _upstream_headers_for_anthropic_route( return MappingProxyType({**caller_headers, **(proxy_auth_header or {})}) +def _upstream_headers_for_bedrock_agent_runtime_route( + request: Request, user_api_key_dict: UserAPIKeyAuth, signed_headers: Mapping[str, object] +) -> Mapping[str, object]: + caller_headers: Final = _caller_headers_without_litellm_secrets( + request, + user_api_key_dict, + _HEADERS_NEVER_FORWARDED_TO_BEDROCK | frozenset(name.lower() for name in signed_headers), + ) + return MappingProxyType({**caller_headers, **signed_headers}) + + async def _prepare_vertex_auth_headers( request: Request, vertex_credentials: VertexPassThroughCredentials | None, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index e0785b002b2..7f044486e14 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1983,6 +1983,146 @@ class TestBedrockAgentRuntimePassthroughToggle: create_route.assert_called_once() +class TestBedrockAgentRuntimePassthroughVirtualKeyLeak: + """Regression for LIT-7912: the agent-runtime branch of ``/bedrock/{endpoint}`` forwarded every caller header, + so a LiteLLM key presented in ``x-api-key`` or ``x-litellm-api-key`` rode to AWS next to the SigV4 signature.""" + + VKEY: Final = "sk-litellm-victim-key" + MASTER_KEY: Final = "sk-master-1234" + ENDPOINT: Final = "knowledgebases/KB1234567/retrieve" + AMBIENT_AWS_ENV: Final = ( + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_SESSION_TOKEN", + "AWS_SESSION_NAME", + "AWS_PROFILE_NAME", + "AWS_ROLE_NAME", + "AWS_WEB_IDENTITY_TOKEN", + "AWS_STS_ENDPOINT", + "AWS_EXTERNAL_ID", + ) + + async def _upstream_headers(self, monkeypatch, headers: list[tuple[bytes, bytes]]) -> dict: + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import HttpPassThroughEndpointHelpers + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", self.MASTER_KEY) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + for ambient in self.AMBIENT_AWS_ENV: + monkeypatch.delenv(ambient, raising=False) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "ak") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "sk") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + caller: Final = UserAPIKeyAuth(api_key=self.VKEY) + + async def receive(): + return {"type": "http.request", "body": b'{"retrievalQuery": {"text": "hi"}}', "more_body": False} + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": f"/bedrock/{self.ENDPOINT}", + "headers": headers, + "query_string": b"", + }, + receive=receive, + ) + captured: dict = {} + + def fake_create_pass_through_route(**kwargs): + captured.update(kwargs) + return AsyncMock(return_value={"status": "success"}) + + module: Final = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints" + with ( + patch(f"{module}.create_request_copy", Mock()), + patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route), + ): + await bedrock_proxy_route( + endpoint=self.ENDPOINT, + request=request, + fastapi_response=Response(), + user_api_key_dict=caller, + ) + return HttpPassThroughEndpointHelpers.forward_headers_from_request( + request_headers=dict(request.headers), + headers=dict(captured["custom_headers"] or {}), + forward_headers=captured.get("_forward_headers", False), + ) + + @staticmethod + def _blob(upstream: dict) -> str: + return " ".join(f"{name}:{value}" for name, value in upstream.items()) + + @staticmethod + def _names_matching(upstream: dict, lowercase_name: str) -> list[str]: + return [name for name in upstream if name.lower() == lowercase_name] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "header_name", ["x-api-key", "x-litellm-api-key", "api-key", "x-goog-api-key", "ocp-apim-subscription-key"] + ) + async def test_virtual_key_in_a_credential_header_never_reaches_aws(self, monkeypatch, header_name: str): + upstream: Final = await self._upstream_headers( + monkeypatch, + [ + (header_name.encode(), self.VKEY.encode()), + (b"content-type", b"application/json"), + (b"x-request-id", b"trace-1"), + ], + ) + + assert self.VKEY not in self._blob(upstream) + assert self._names_matching(upstream, header_name) == [] + assert upstream["x-request-id"] == "trace-1", "a benign caller header still reaches AWS" + assert upstream["Authorization"].startswith("AWS4-HMAC-SHA256") + assert self._names_matching(upstream, "content-type") == ["Content-Type"], "the signed header is the only one" + + @pytest.mark.asyncio + async def test_credential_headers_are_dropped_by_name_even_when_they_carry_someone_elses_key(self, monkeypatch): + other_key: Final = "sk-other-tenant-key" + upstream: Final = await self._upstream_headers( + monkeypatch, + [ + (b"x-api-key", other_key.encode()), + (b"x-litellm-api-key", other_key.encode()), + (b"x-request-id", b"trace-3"), + ], + ) + + assert other_key not in self._blob(upstream) + assert self._names_matching(upstream, "x-api-key") == [] + assert self._names_matching(upstream, "x-litellm-api-key") == [] + assert upstream["x-request-id"] == "trace-3" + + @pytest.mark.asyncio + async def test_virtual_key_in_authorization_bearer_is_replaced_by_the_sigv4_signature(self, monkeypatch): + upstream: Final = await self._upstream_headers( + monkeypatch, + [(b"authorization", f"Bearer {self.VKEY}".encode()), (b"content-type", b"application/json")], + ) + + assert self.VKEY not in self._blob(upstream) + assert self._names_matching(upstream, "authorization") == ["Authorization"] + assert upstream["Authorization"].startswith("AWS4-HMAC-SHA256") + + @pytest.mark.asyncio + async def test_authenticated_secrets_in_any_other_header_never_reach_aws(self, monkeypatch): + upstream: Final = await self._upstream_headers( + monkeypatch, + [ + (b"x-api-key", self.VKEY.encode()), + (b"x-forwarded-key", self.VKEY.encode()), + (b"x-operator-token", self.MASTER_KEY.encode()), + (b"x-request-id", b"trace-2"), + ], + ) + + assert self.VKEY not in self._blob(upstream) and self.MASTER_KEY not in self._blob(upstream) + assert self._names_matching(upstream, "x-forwarded-key") == [] + assert self._names_matching(upstream, "x-operator-token") == [] + assert upstream["x-request-id"] == "trace-2" + + class TestLLMPassthroughFactoryProxyRoute: @pytest.mark.asyncio async def test_llm_passthrough_factory_proxy_route_success(self): From 610249c66341ce69d68b2e70b4d35b3df9934cef Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:08:46 -0700 Subject: [PATCH 162/187] fix(bedrock): carry the tool_call_id into the rewritten tool call text --- litellm/llms/bedrock/chat/converse_transformation.py | 4 +++- .../llms/bedrock/chat/test_converse_transformation.py | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 33a77e8d346..7c2dadbc970 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -228,7 +228,9 @@ class AmazonConverseConfig(BaseConfig): function = tool_call.get("function") or {} name = function.get("name") or "unknown_tool" arguments = function.get("arguments") or "" - return f"[tool call: {name}({arguments})]" + call_id = tool_call.get("id") + label = f"tool call {call_id}" if call_id else "tool call" + return f"[{label}: {name}({arguments})]" def _result_text(message: AllMessageValues) -> str: rendered = convert_content_list_to_str(message).strip() diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 21e8e94d871..9ecca87134b 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -6467,7 +6467,8 @@ def test_neutralize_orphaned_tool_blocks_rewrites_when_no_tools(): # '{"city": "Paris"}' is escaped, so assert on quote-free tokens that survive. assert "city" in serialized and "Paris" in serialized assert "Sunny, 25C" in serialized - assert "call_abc" in serialized # tool_call_id correlation preserved + assert "[tool call call_abc: get_weather(" in result[1]["content"] + assert "[tool result for call_abc: Sunny, 25C]" in result[2]["content"] @pytest.mark.parametrize("tools_value", [[], None]) From 81ae5caa7e65e05be6a07f53fa46b394e8003462 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 16:11:52 -0700 Subject: [PATCH 163/187] fix(proxy): pass only a team admin's changed fields on to the team update --- .../team_admin_field_permissions.py | 21 ++++++++--- .../management_endpoints/team_endpoints.py | 4 +-- .../test_team_admin_field_permissions.py | 33 +++++++++++------ .../test_team_endpoints.py | 36 +++++++++++++++++++ 4 files changed, 77 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py index 6e836ad440c..4248501551f 100644 --- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -35,6 +35,7 @@ _SETTINGS_LOCATION: Final = "Settings > UI > Team admin editable fields" @dataclass(frozen=True, slots=True) class TeamAdminEditAllowed: + request: UpdateTeamRequest kind: Literal["allowed"] = "allowed" @@ -143,6 +144,15 @@ def changed_team_fields(data: UpdateTeamRequest, existing_row: LiteLLM_TeamTable return column_changes | _metadata_changes(data, submitted, existing) +def _only_changes(data: UpdateTeamRequest, changed: frozenset[str]) -> UpdateTeamRequest: + """The request without the values it resends unchanged, which would otherwise still trigger derived writes + such as a resent budget_duration pushing budget_reset_at back.""" + sent: Final = frozenset(data.model_fields_set) + via_metadata: Final = frozenset({"metadata"}) if changed - sent else frozenset() + kept: Final = frozenset({"team_id"}) | (changed & sent) | via_metadata + return UpdateTeamRequest.model_validate(data.model_dump(include=MappingProxyType({field: True for field in kept}))) + + def team_admin_edit_verdict( data: UpdateTeamRequest, existing: LiteLLM_TeamTable, @@ -150,16 +160,17 @@ def team_admin_edit_verdict( ) -> TeamAdminEditVerdict: if not permitted: return TeamAdminEditingDisabled() - blocked: Final = sorted(changed_team_fields(data, existing) - permitted) + changed: Final = changed_team_fields(data, existing) + blocked: Final = sorted(changed - permitted) if blocked: return TeamAdminFieldNotPermitted(field=blocked[0]) - return TeamAdminEditAllowed() + return TeamAdminEditAllowed(request=_only_changes(data, changed)) -def raise_for_team_admin_edit_verdict(verdict: TeamAdminEditVerdict) -> None: +def team_admin_request_or_raise(verdict: TeamAdminEditVerdict) -> UpdateTeamRequest: match verdict: - case TeamAdminEditAllowed(): - return + case TeamAdminEditAllowed(request=request): + return request case TeamAdminEditingDisabled(): raise HTTPException( status_code=403, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 72367b2bced..b2dc3551ced 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -140,9 +140,9 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import ( ) from litellm.proxy.management_endpoints.team_admin_field_permissions import ( SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS, - raise_for_team_admin_edit_verdict, resolve_team_admin_editable_fields, team_admin_edit_verdict, + team_admin_request_or_raise, ) from litellm.proxy.management_helpers.access_group_team_sync import ( TEAM_ADVISORY_LOCK_SQL, @@ -2218,7 +2218,7 @@ async def update_team( if access_role is None: _raise_team_access_denied() if access_role == "team_admin": - raise_for_team_admin_edit_verdict( + data = team_admin_request_or_raise( # rebind-ok: resent values must not reach the derived writes below team_admin_edit_verdict( data=data, existing=existing_team, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py index 91479921c61..5b31089f91e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -7,9 +7,9 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import ( TeamAdminEditingDisabled, TeamAdminFieldNotPermitted, changed_team_fields, - raise_for_team_admin_edit_verdict, resolve_team_admin_editable_fields, team_admin_edit_verdict, + team_admin_request_or_raise, ) _SUPPORTED = frozenset({"tpm_limit", "rpm_limit", "team_alias"}) @@ -96,10 +96,22 @@ class TestTeamAdminEditVerdict: verdict = team_admin_edit_verdict(UpdateTeamRequest(team_id="team-1"), _team(), frozenset()) assert verdict == TeamAdminEditingDisabled() - def test_changes_within_permitted_fields_are_allowed(self): - data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha") - verdict = team_admin_edit_verdict(data, _team(team_alias="alpha"), frozenset({"tpm_limit"})) - assert verdict == TeamAdminEditAllowed() + def test_allowed_request_keeps_only_the_changed_fields(self): + data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha", budget_duration="30d") + existing = _team(team_alias="alpha", budget_duration="30d") + verdict = team_admin_edit_verdict(data, existing, frozenset({"tpm_limit"})) + assert isinstance(verdict, TeamAdminEditAllowed) + assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6} + + def test_permitted_field_changed_inside_metadata_keeps_the_metadata(self): + data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["b"]}, team_alias="alpha") + existing = _team(team_alias="alpha", metadata={"guardrails": ["a"]}) + verdict = team_admin_edit_verdict(data, existing, frozenset({"guardrails"})) + assert isinstance(verdict, TeamAdminEditAllowed) + assert verdict.request.model_dump(exclude_unset=True) == { + "team_id": "team-1", + "metadata": {"guardrails": ["b"]}, + } def test_first_blocked_field_in_sorted_order_is_reported(self): data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, rpm_limit=6, blocked=True) @@ -107,19 +119,20 @@ class TestTeamAdminEditVerdict: assert verdict == TeamAdminFieldNotPermitted(field="blocked") -class TestRaiseForTeamAdminEditVerdict: - def test_allowed_does_not_raise(self): - assert raise_for_team_admin_edit_verdict(TeamAdminEditAllowed()) is None +class TestTeamAdminRequestOrRaise: + def test_allowed_hands_back_its_request(self): + request = UpdateTeamRequest(team_id="team-1", tpm_limit=6) + assert team_admin_request_or_raise(TeamAdminEditAllowed(request=request)) is request def test_disabled_is_a_403_pointing_at_the_proxy_admin(self): with pytest.raises(HTTPException) as exc: - raise_for_team_admin_edit_verdict(TeamAdminEditingDisabled()) + team_admin_request_or_raise(TeamAdminEditingDisabled()) assert exc.value.status_code == 403 assert "cannot edit team settings" in exc.value.detail assert "Settings > UI > Team admin editable fields" in exc.value.detail def test_field_not_permitted_is_a_403_naming_the_field(self): with pytest.raises(HTTPException) as exc: - raise_for_team_admin_edit_verdict(TeamAdminFieldNotPermitted(field="blocked")) + team_admin_request_or_raise(TeamAdminFieldNotPermitted(field="blocked")) assert exc.value.status_code == 403 assert "'blocked'" in exc.value.detail diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index f507311a24f..3dfd994bcee 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -15078,6 +15078,42 @@ async def test_update_team_team_admin_changes_tpm_limit_once_a_proxy_admin_enabl assert "'rpm_limit'" in str(refused.value.message) +@pytest.mark.asyncio +async def test_update_team_team_admin_resending_budget_settings_does_not_push_back_budget_resets( + disable_audit_logging_for_mocked_team, +): + """A resent budget_duration or budget_limits would otherwise recompute the reset timestamps from now.""" + import contextlib + + stored_windows = [{"budget_duration": "7d", "max_budget": 5.0, "reset_at": "2026-09-20T00:00:00Z"}] + budgeted_team = MagicMock() + budgeted_team.metadata = {} + budgeted_team.model_dump.return_value = { + "team_id": "test_team_id", + "team_alias": "test_team", + "metadata": {}, + "budget_duration": "30d", + "budget_limits": stored_windows, + "members_with_roles": [{"user_id": "team-admin", "role": "admin"}], + } + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {}) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=budgeted_team) + stack.enter_context(_team_admin_may_edit("tpm_limit")) + await update_team( + data=UpdateTeamRequest( + team_id="test_team_id", tpm_limit=5000, budget_duration="30d", budget_limits=stored_windows + ), + http_request=_update_request_stub(), + user_api_key_dict=_TEAM_ADMIN_CALLER, + ) + + written = prisma.db.litellm_teamtable.update.call_args.kwargs["data"] + assert written["tpm_limit"] == 5000 + assert not {"budget_duration", "budget_reset_at", "budget_limits"} & written.keys() + + @pytest.mark.asyncio async def test_update_team_holds_a_team_admin_to_the_org_tpm_limit(disable_audit_logging_for_mocked_team): """The org ceiling lives on the org's budget row, so /team/update must load it to enforce the cap.""" From 417cc5c4fb61dc244cb5da905a3b53afc7ff65d6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 09:58:14 -0700 Subject: [PATCH 164/187] feat(ui): keep organizations list and detail tab state in the URL The organizations list now reads its search (org_search), org ID filter (filter_org_id), sort (sort_by, sort_order) and pagination (page, page_size) from the URL through useUrlTableState. The detail view tabs are controlled by ?org_tab=, and the Edit row action opens ?org=&org_tab=settings in one history entry instead of passing an editOrg flag --- .../_components/OrganizationsPanel.test.tsx | 109 ++++++++++++-- .../_components/OrganizationsPanel.tsx | 47 +++--- .../_components/OrganizationsTable.test.tsx | 134 +++++++++++++++--- .../_components/OrganizationsTable.tsx | 12 +- .../_components/useOrganizationsTableState.ts | 19 +++ .../organization/organizationTabs.ts | 3 + .../organization/organization_view.test.tsx | 112 +++++++++++++-- .../organization/organization_view.tsx | 14 +- 8 files changed, 380 insertions(+), 70 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/useOrganizationsTableState.ts create mode 100644 ui/litellm-dashboard/src/components/organization/organizationTabs.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx index c15b9fcaddb..3e87492778a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx @@ -1,10 +1,23 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { act, render, screen, waitFor } from "@testing-library/react"; +import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; import { NuqsTestingAdapter, type UrlUpdateEvent } from "nuqs/adapters/testing"; import React from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import type OrganizationsTableComponent from "./OrganizationsTable"; import type OrganizationInfoViewComponent from "@/components/organization/organization_view"; +import type { OrganizationListFilters } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; + +const useOrganizationsSpy = vi.hoisted(() => vi.fn<(filters?: OrganizationListFilters) => void>()); +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + useOrganizations: (filters?: OrganizationListFilters) => { + useOrganizationsSpy(filters); + return actual.useOrganizations(filters); + }, + }; +}); vi.mock("@/components/vector_store_management/VectorStoreSelector", () => ({ __esModule: true, @@ -79,10 +92,13 @@ const renderPanel = ({ premiumUser = true, searchParams = "" }: RenderPanelOptio const expectQueryString = (queryString: string) => waitFor(() => expect(onUrlUpdate).toHaveBeenLastCalledWith(expect.objectContaining({ queryString }))); +const lastSearchParams = () => onUrlUpdate.mock.calls.at(-1)?.[0].searchParams; + beforeEach(() => { capturedTableProps = null; mockOrgInfoView.mockClear(); onUrlUpdate.mockClear(); + useOrganizationsSpy.mockClear(); }); describe("OrganizationsPanel", () => { @@ -123,9 +139,7 @@ describe("OrganizationsPanel - org detail deep link (?org=)", () => { it("opens the org detail directly from a ?org= deep link", () => { renderPanel({ searchParams: "?org=org-from-url" }); - expect(mockOrgInfoView).toHaveBeenLastCalledWith( - expect.objectContaining({ organizationId: "org-from-url", editOrg: false }), - ); + expect(mockOrgInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ organizationId: "org-from-url" })); expect(screen.queryByTestId("organizations-table")).not.toBeInTheDocument(); }); @@ -139,23 +153,24 @@ describe("OrganizationsPanel - org detail deep link (?org=)", () => { expect(screen.getByTestId("organizations-table")).toBeInTheDocument(); }); - it("the edit action opens the detail in edit mode with ?org= set", async () => { + it("the edit action pushes ?org= with ?org_tab=settings in one history entry", async () => { renderPanel(); act(() => capturedTableProps?.onEditClick("org-edit")); - await expectQueryString("?org=org-edit"); - expect(mockOrgInfoView).toHaveBeenLastCalledWith( - expect.objectContaining({ organizationId: "org-edit", editOrg: true }), + await expectQueryString("?org=org-edit&org_tab=settings"); + expect(onUrlUpdate).toHaveBeenCalledTimes(1); + expect(onUrlUpdate).toHaveBeenLastCalledWith( + expect.objectContaining({ options: expect.objectContaining({ history: "push" }) }), ); + expect(mockOrgInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ organizationId: "org-edit" })); }); - it("a plain row click after leaving an edit view via browser history does not reopen in edit mode", async () => { + it("a plain row click after leaving an edit view via browser history opens the detail without the settings tab", async () => { const { navigate } = renderPanel(); act(() => capturedTableProps?.onEditClick("org-edit")); - await expectQueryString("?org=org-edit"); - expect(mockOrgInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ editOrg: true })); + await expectQueryString("?org=org-edit&org_tab=settings"); navigate(""); expect(screen.getByTestId("organizations-table")).toBeInTheDocument(); @@ -163,8 +178,76 @@ describe("OrganizationsPanel - org detail deep link (?org=)", () => { act(() => capturedTableProps?.onOrganizationClick("org-plain")); await expectQueryString("?org=org-plain"); - expect(mockOrgInfoView).toHaveBeenLastCalledWith( - expect.objectContaining({ organizationId: "org-plain", editOrg: false }), + expect(mockOrgInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ organizationId: "org-plain" })); + }); + + it("a row click drops a leftover ?org_tab= so the detail opens on its default tab", async () => { + renderPanel({ searchParams: "?org_tab=settings" }); + + act(() => capturedTableProps?.onOrganizationClick("org-plain")); + + await expectQueryString("?org=org-plain"); + }); + + it("closing the org detail drops ?org_tab= together with ?org=", async () => { + renderPanel({ searchParams: "?org=org-from-url&org_tab=members" }); + + act(() => mockOrgInfoView.mock.calls.at(-1)?.[0].onClose()); + + await expectQueryString(""); + expect(onUrlUpdate).toHaveBeenCalledTimes(1); + expect(onUrlUpdate).toHaveBeenLastCalledWith( + expect.objectContaining({ options: expect.objectContaining({ history: "push" }) }), ); }); }); + +describe("OrganizationsPanel - list filters in the URL", () => { + it("restores the name search and org ID filter from the URL and fetches with both", () => { + renderPanel({ searchParams: "?org_search=Acme&filter_org_id=org-7" }); + + expect(screen.getByPlaceholderText("Search by Organization Name")).toHaveValue("Acme"); + expect(screen.getByPlaceholderText("Search by Organization ID")).toHaveValue("org-7"); + expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "org-7", org_alias: "Acme" }); + expect(capturedTableProps?.searchActive).toBe(true); + }); + + it("keeps the org ID filter panel collapsed when the URL has no org ID filter", () => { + renderPanel({ searchParams: "?org_search=Acme" }); + + expect(screen.queryByPlaceholderText("Search by Organization ID")).not.toBeInTheDocument(); + expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "", org_alias: "Acme" }); + }); + + it("writes the name search to ?org_search= and returns the list to the first page", async () => { + renderPanel({ searchParams: "?page=3" }); + + fireEvent.change(screen.getByPlaceholderText("Search by Organization Name"), { target: { value: "Acme" } }); + + await waitFor(() => expect(lastSearchParams()?.get("org_search")).toBe("Acme")); + expect(lastSearchParams()?.has("page")).toBe(false); + expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "", org_alias: "Acme" }); + }); + + it("writes the org ID filter to ?filter_org_id= and returns the list to the first page", async () => { + renderPanel({ searchParams: "?page=3" }); + + fireEvent.click(screen.getByRole("button", { name: "Filters" })); + fireEvent.change(screen.getByPlaceholderText("Search by Organization ID"), { target: { value: "org-9" } }); + + await waitFor(() => expect(lastSearchParams()?.get("filter_org_id")).toBe("org-9")); + expect(lastSearchParams()?.has("page")).toBe(false); + expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "org-9", org_alias: "" }); + }); + + it("clears the search, the org ID filter and the page in one update on reset", async () => { + renderPanel({ searchParams: "?org_search=Acme&filter_org_id=org-7&page=2" }); + + fireEvent.click(screen.getByRole("button", { name: "Reset Filters" })); + + await expectQueryString(""); + expect(onUrlUpdate).toHaveBeenCalledTimes(1); + expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "", org_alias: "" }); + expect(capturedTableProps?.searchActive).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.tsx index 4fe4cf47b9c..2810902bef6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.tsx @@ -2,16 +2,18 @@ import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/orga import { useUserModels } from "@/app/(dashboard)/hooks/models/useModels"; import OrganizationFilters, { FilterState } from "@/app/(dashboard)/organizations/OrganizationFilters"; import { useQueryClient } from "@tanstack/react-query"; -import { parseAsString, useQueryState } from "nuqs"; +import { parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs"; import React, { useState } from "react"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; import { toast } from "@/lib/toast"; import { organizationDeleteCall } from "@/components/networking"; import { OrgCreateDialog } from "@/components/organization/org-create/OrgCreateDialog"; import OrganizationInfoView from "@/components/organization/organization_view"; +import { ORGANIZATION_TAB_URL_KEY, ORGANIZATION_TABS } from "@/components/organization/organizationTabs"; import { Button } from "@/components/ui/button"; import OrganizationsTable from "./OrganizationsTable"; +import { organizationIdFilter, useOrganizationsTableState } from "./useOrganizationsTableState"; interface OrganizationsPanelProps { userRole: string; @@ -19,15 +21,25 @@ interface OrganizationsPanelProps { premiumUser: boolean; } +const ORGANIZATION_DETAIL_STATE = { + org: parseAsString, + tab: parseAsStringLiteral(ORGANIZATION_TABS), +}; +const ORGANIZATION_DETAIL_URL_KEYS = { tab: ORGANIZATION_TAB_URL_KEY }; + const OrganizationsPanel: React.FC = ({ userRole, accessToken, premiumUser }) => { - const [selectedOrgId, setSelectedOrgId] = useQueryState("org", parseAsString.withOptions({ history: "push" })); - const [editOrg, setEditOrg] = useState(false); + const [{ org: selectedOrgId }, setOrganizationDetail] = useQueryStates(ORGANIZATION_DETAIL_STATE, { + history: "push", + urlKeys: ORGANIZATION_DETAIL_URL_KEYS, + }); + const tableState = useOrganizationsTableState(); + const { setSearch, onColumnFiltersChange } = tableState; + const filters: FilterState = { org_id: organizationIdFilter(tableState), org_alias: tableState.search }; const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [orgToDelete, setOrgToDelete] = useState(null); const [isDeleting, setIsDeleting] = useState(false); const [isOrgModalVisible, setIsOrgModalVisible] = useState(false); - const [showFilters, setShowFilters] = useState(false); - const [filters, setFilters] = useState({ org_id: "", org_alias: "" }); + const [showFilters, setShowFilters] = useState(() => filters.org_id !== ""); const queryClient = useQueryClient(); const { data: organizations = [], isLoading } = useOrganizations({ @@ -41,11 +53,16 @@ const OrganizationsPanel: React.FC = ({ userRole, acces const refetchOrganizations = () => queryClient.invalidateQueries({ queryKey: organizationKeys.lists() }); const handleFilterChange = (key: keyof FilterState, value: string) => { - setFilters((previousFilters) => ({ ...previousFilters, [key]: value })); + if (key === "org_alias") { + setSearch(value); + return; + } + onColumnFiltersChange(value ? [{ id: "org_id", value }] : []); }; const handleFilterReset = () => { - setFilters({ org_id: "", org_alias: "" }); + setSearch(""); + onColumnFiltersChange([]); }; const handleDelete = (orgId: string | null) => { @@ -108,15 +125,11 @@ const OrganizationsPanel: React.FC = ({ userRole, acces {selectedOrgId ? ( { - void setSelectedOrgId(null); - setEditOrg(false); - }} + onClose={() => void setOrganizationDetail(null)} accessToken={accessToken} is_org_admin={true} is_proxy_admin={userRole === "Admin"} userModels={userModels} - editOrg={editOrg} /> ) : ( <> @@ -133,14 +146,8 @@ const OrganizationsPanel: React.FC = ({ userRole, acces isLoading={isLoading} userRole={userRole} searchActive={searchActive} - onOrganizationClick={(organizationId) => { - setEditOrg(false); - void setSelectedOrgId(organizationId); - }} - onEditClick={(organizationId) => { - void setSelectedOrgId(organizationId); - setEditOrg(true); - }} + onOrganizationClick={(organizationId) => void setOrganizationDetail({ org: organizationId, tab: null })} + onEditClick={(organizationId) => void setOrganizationDetail({ org: organizationId, tab: "settings" })} onDeleteClick={handleDelete} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx index 4bf465b847b..9d163fe2c08 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx @@ -1,7 +1,9 @@ -import { render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; import React from "react"; -import { describe, expect, it, vi } from "vitest"; +import { describe, expect, it, vi, type Mock } from "vitest"; + +import { renderWithProviders, screen, waitFor, within } from "../../../../../tests/test-utils"; import { Organization } from "@/components/networking"; @@ -26,6 +28,34 @@ const makeOrganization = (overrides: Partial = {}): Organization = ...overrides, }); +const thirtyOrganizations = Array.from({ length: 30 }, (_, index) => + makeOrganization({ organization_id: `org-${index}`, organization_alias: `Org ${index}` }), +); + +const sortableOrganization = (alias: string, createdAt: string, spend: number): Organization => { + const overrides: Partial = { + organization_id: `org-${alias.toLowerCase()}`, + organization_alias: alias, + created_at: createdAt, + spend, + }; + return makeOrganization(overrides); +}; + +const sortableOrganizations = [ + sortableOrganization("Mid", "2024-03-01T00:00:00Z", 5), + sortableOrganization("Zed", "2023-01-01T00:00:00Z", 1), + sortableOrganization("Ace", "2025-01-01T00:00:00Z", 3), +]; + +const bodyRowAliases = () => + screen + .getAllByRole("row") + .slice(1) + .map((row) => ["Ace", "Mid", "Zed"].find((alias) => within(row).queryByText(alias) !== null)); + +const lastSearchParams = (onUrlUpdate: Mock) => onUrlUpdate.mock.calls.at(-1)?.[0].searchParams; + const baseProps = { isLoading: false, userRole: "Admin", @@ -37,7 +67,7 @@ const baseProps = { describe("OrganizationsTable", () => { it("renders every column header", () => { - render(); + renderWithProviders(); for (const header of [ "Organization ID", "Organization Name", @@ -55,7 +85,7 @@ describe("OrganizationsTable", () => { it("opens the detail view when the organization ID cell is clicked", async () => { const user = userEvent.setup(); const onOrganizationClick = vi.fn(); - render( + renderWithProviders( { const user = userEvent.setup(); const onEditClick = vi.fn(); const onDeleteClick = vi.fn(); - render( + renderWithProviders( { }); it("hides the row actions menu from non-admins", () => { - render( + renderWithProviders( { }); it("sorts by created_at descending by default", () => { - render( + renderWithProviders( { }); it("renders budget, limits, members, and models for a fully-populated organization", () => { - render( + renderWithProviders( { }); it("shows Unlimited budget and All Proxy Models when unset", () => { - render( + renderWithProviders( { }); it("renders a tpm/rpm limit of 0 as 0, never as Unlimited", () => { - render( + renderWithProviders( { }); it("renders loading skeletons instead of rows while loading", () => { - render( + renderWithProviders( { it("pages long lists client-side with the shared size selector and footer", async () => { const user = userEvent.setup(); - const organizations = Array.from({ length: 30 }, (_, index) => - makeOrganization({ organization_id: `org-${index}`, organization_alias: `Org ${index}` }), - ); - render(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); expect(screen.getAllByRole("row")).toHaveLength(26); expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 30"); @@ -207,13 +235,87 @@ describe("OrganizationsTable", () => { expect(screen.getAllByRole("row")).toHaveLength(31); expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-30 of 30"); + await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("page_size")).toBe("50")); }); it("uses a search-aware empty state", () => { - const { rerender } = render(); + const { rerender } = renderWithProviders( + , + ); expect(screen.getByText("No organizations yet")).toBeInTheDocument(); rerender(); expect(screen.getByText("No matching organizations")).toBeInTheDocument(); }); }); + +describe("OrganizationsTable URL state", () => { + it("restores the sort column and direction from ?sort_by=&sort_order=", () => { + renderWithProviders(, { + searchParams: "?sort_by=spend&sort_order=desc", + }); + + expect(bodyRowAliases()).toEqual(["Mid", "Ace", "Zed"]); + }); + + it("falls back to sorting by creation date for a ?sort_by= column that cannot be sorted", () => { + renderWithProviders(, { + searchParams: "?sort_by=members&sort_order=asc", + }); + + expect(bodyRowAliases()).toEqual(["Zed", "Mid", "Ace"]); + }); + + it("writes the clicked sort column to the URL and returns to the first page", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { + searchParams: "?page=2", + onUrlUpdate, + }); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 26-30 of 30"); + + await user.click(screen.getByTestId("sort-header-organization_alias")); + + await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("sort_by")).toBe("organization_alias")); + expect(onUrlUpdate).toHaveBeenCalledTimes(1); + expect(lastSearchParams(onUrlUpdate)?.get("sort_order")).toBe("asc"); + expect(lastSearchParams(onUrlUpdate)?.has("page")).toBe(false); + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 30"); + expect(within(screen.getAllByRole("row")[1]).getByText("Org 0")).toBeInTheDocument(); + }); + + it("opens the page named by ?page= and writes page changes back to the URL", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { + searchParams: "?page=2", + onUrlUpdate, + }); + + expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 26-30 of 30"); + expect(screen.getByText("org-29")).toBeInTheDocument(); + + await user.click(screen.getByTestId("pagination-prev")); + + await waitFor(() => expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 30")); + expect(lastSearchParams(onUrlUpdate)?.has("page")).toBe(false); + + await user.click(screen.getByTestId("pagination-next")); + + await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("page")).toBe("2")); + }); + + it("keeps a deep-linked ?page= while the organization list is still loading", async () => { + const onUrlUpdate = vi.fn(); + const { rerender } = renderWithProviders(, { + searchParams: "?page=2", + onUrlUpdate, + }); + + rerender(); + + await waitFor(() => expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 26-30 of 30")); + expect(onUrlUpdate).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx index dbf516d75ae..a9ac0b7e699 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx @@ -1,13 +1,13 @@ "use client"; -import { SortingState } from "@tanstack/react-table"; import { Building2, SearchX } from "lucide-react"; -import React, { useMemo, useState } from "react"; +import React, { useMemo } from "react"; import { DataTable } from "@/components/shared/DataTable"; import { Organization } from "@/components/networking"; import { getOrganizationsTableColumns } from "./OrganizationsTableColumns"; +import { useOrganizationsTableState } from "./useOrganizationsTableState"; interface OrganizationsTableProps { organizations: Organization[]; @@ -19,8 +19,6 @@ interface OrganizationsTableProps { onDeleteClick: (organizationId: string) => void; } -const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; - function EmptyState({ searchActive }: { searchActive: boolean }) { const Icon = searchActive ? SearchX : Building2; return ( @@ -49,7 +47,7 @@ const OrganizationsTable: React.FC = ({ onEditClick, onDeleteClick, }) => { - const [sorting, setSorting] = useState(DEFAULT_SORTING); + const { sorting, onSortingChange, pagination, onPaginationChange } = useOrganizationsTableState(); const columns = useMemo(() => { const deps = { userRole, onOrganizationClick, onEditClick, onDeleteClick }; @@ -60,11 +58,13 @@ const OrganizationsTable: React.FC = ({ organization.organization_id || String(index)} sortingMode="client" sorting={sorting} - onSortingChange={setSorting} + onSortingChange={onSortingChange} isLoading={isLoading} loadingMessage="Loading organizations…" noDataMessage={} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/useOrganizationsTableState.ts b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/useOrganizationsTableState.ts new file mode 100644 index 00000000000..20a54a25ba1 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/useOrganizationsTableState.ts @@ -0,0 +1,19 @@ +import { useUrlTableState, type UrlTableState, type UrlTableStateOptions } from "@/components/shared/DataTable"; + +const FILTER_COLUMNS = ["org_id"] as const; +type FilterColumn = (typeof FILTER_COLUMNS)[number]; + +const TABLE_STATE_OPTIONS: UrlTableStateOptions = { + sortFields: ["organization_id", "organization_alias", "created_at", "spend"], + defaultSort: { id: "created_at", desc: true }, + defaultPageSize: 25, + filterColumns: FILTER_COLUMNS, + urlKeys: { search: "org_search" }, +}; + +export const useOrganizationsTableState = (): UrlTableState => useUrlTableState(TABLE_STATE_OPTIONS); + +export const organizationIdFilter = ({ columnFilters }: Pick): string => { + const value = columnFilters.find((filter) => filter.id === "org_id")?.value; + return typeof value === "string" ? value : ""; +}; diff --git a/ui/litellm-dashboard/src/components/organization/organizationTabs.ts b/ui/litellm-dashboard/src/components/organization/organizationTabs.ts new file mode 100644 index 00000000000..db0a2692356 --- /dev/null +++ b/ui/litellm-dashboard/src/components/organization/organizationTabs.ts @@ -0,0 +1,3 @@ +export const ORGANIZATION_TABS = ["overview", "members", "settings"] as const; +export type OrganizationTab = (typeof ORGANIZATION_TABS)[number]; +export const ORGANIZATION_TAB_URL_KEY = "org_tab"; diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx index 799cd1adffe..d609b657717 100644 --- a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx +++ b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx @@ -1,8 +1,10 @@ import React from "react"; -import { fireEvent, screen, waitFor, within } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { vi, test, expect, beforeEach } from "vitest"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { QueryClientProvider } from "@tanstack/react-query"; +import { NuqsTestingAdapter, type OnUrlUpdateFunction } from "nuqs/adapters/testing"; +import { vi, test, expect, beforeEach, describe, type Mock } from "vitest"; +import { renderWithProviders, testQueryClient } from "../../../tests/test-utils"; import OrganizationInfoView from "./organization_view"; import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; @@ -115,7 +117,6 @@ test("renders organization view after loading data", async () => { is_org_admin={false} is_proxy_admin={false} userModels={[]} - editOrg={false} />, ); @@ -135,7 +136,6 @@ test("should display empty state when organization has no members", async () => is_org_admin={false} is_proxy_admin={false} userModels={[]} - editOrg={false} />, ); @@ -165,7 +165,6 @@ test("should display team aliases when teams are available", async () => { is_org_admin={false} is_proxy_admin={false} userModels={[]} - editOrg={false} />, ); @@ -199,7 +198,6 @@ test("should display team ID as fallback when alias is not found", async () => { is_org_admin={false} is_proxy_admin={false} userModels={[]} - editOrg={false} />, ); @@ -223,7 +221,6 @@ test("links each team badge to that team's detail page", async () => { is_org_admin={false} is_proxy_admin={false} userModels={[]} - editOrg={false} />, ); @@ -250,7 +247,6 @@ test("model badges stay non-clickable", async () => { is_org_admin={false} is_proxy_admin={false} userModels={[]} - editOrg={false} />, ); @@ -272,7 +268,6 @@ test("should keep unsaved settings edits when switching tabs and back", async () is_org_admin={false} is_proxy_admin={true} userModels={[]} - editOrg={false} />, ); @@ -308,7 +303,6 @@ test("renders a tpm/rpm limit of 0 as 0 in the overview and settings tabs, never is_org_admin={false} is_proxy_admin={true} userModels={[]} - editOrg={false} />, ); @@ -323,3 +317,99 @@ test("renders a tpm/rpm limit of 0 as 0 in the overview and settings tabs, never expect(screen.queryByText("TPM: Unlimited")).not.toBeInTheDocument(); expect(screen.queryByText("RPM: Unlimited")).not.toBeInTheDocument(); }); + +const renderOrgView = (props: { is_proxy_admin?: boolean } = {}) => ( + {}} + accessToken="test-token" + is_org_admin={false} + is_proxy_admin={props.is_proxy_admin ?? false} + userModels={[]} + /> +); + +const lastSearchParams = (onUrlUpdate: Mock) => onUrlUpdate.mock.calls.at(-1)?.[0].searchParams; + +describe("organization detail tab in the URL (?org_tab=)", () => { + beforeEach(() => { + mockUseOrganization.mockReturnValue({ data: mockOrg, isLoading: false } as unknown as ReturnType< + typeof useOrganization + >); + }); + + test("opens on the tab named by ?org_tab=", () => { + renderWithProviders(renderOrgView(), { searchParams: "?org=org_123&org_tab=members" }); + + expect(screen.getByRole("tab", { name: "Members" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByText("No members found")).toBeInTheDocument(); + }); + + test("the settings deep link used by the list's Edit action opens the Settings tab", () => { + renderWithProviders(renderOrgView({ is_proxy_admin: true }), { searchParams: "?org=org_123&org_tab=settings" }); + + expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument(); + }); + + test("opens on Overview when the URL names no tab", () => { + renderWithProviders(renderOrgView(), { searchParams: "?org=org_123" }); + + expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + }); + + test("writes the selected tab to ?org_tab= and drops it again for Overview", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(renderOrgView(), { searchParams: "?org=org_123", onUrlUpdate }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("org_tab")).toBe("settings")); + expect(lastSearchParams(onUrlUpdate)?.get("org")).toBe("org_123"); + expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "true"); + + await user.click(screen.getByRole("tab", { name: "Overview" })); + + await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.has("org_tab")).toBe(false)); + expect(lastSearchParams(onUrlUpdate)?.get("org")).toBe("org_123"); + }); + + test("falls back to Overview for an unknown ?org_tab= and removes it from the URL", async () => { + const onUrlUpdate = vi.fn(); + render(renderOrgView(), { + wrapper: ({ children }) => ( + + {children} + + ), + }); + + expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + expect(lastSearchParams(onUrlUpdate)?.has("org_tab")).toBe(false); + expect(lastSearchParams(onUrlUpdate)?.get("org")).toBe("org_123"); + }); + + test("follows back and forward navigation between tabs while the detail view stays open", () => { + const atUrl = (searchParams: string) => ( + + {renderOrgView()} + + ); + const { rerender } = render(atUrl("?org=org_123&org_tab=members")); + expect(screen.getByRole("tab", { name: "Members" })).toHaveAttribute("aria-selected", "true"); + + rerender(atUrl("?org=org_123")); + expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + + rerender(atUrl("?org=org_123&org_tab=members")); + expect(screen.getByRole("tab", { name: "Members" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByText("No members found")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.tsx index 096e736b493..c800d12ad62 100644 --- a/ui/litellm-dashboard/src/components/organization/organization_view.tsx +++ b/ui/litellm-dashboard/src/components/organization/organization_view.tsx @@ -1,6 +1,7 @@ import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { organizationKeys, useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useQueryClient } from "@tanstack/react-query"; +import { useUrlTab } from "@/hooks/useUrlTab"; import { useVisitedTabs } from "@/hooks/useVisitedTabs"; import { MoneyCell } from "@/components/shared/table_cells"; import CopyButton from "@/components/shared/CopyButton"; @@ -25,6 +26,7 @@ import { import ObjectPermissionsView from "../object_permissions_view"; import MemberModal from "../team/EditMembership"; import { OrgSettingsForm } from "./org-settings/OrgSettingsForm"; +import { ORGANIZATION_TAB_URL_KEY, ORGANIZATION_TABS, type OrganizationTab } from "./organizationTabs"; interface OrganizationInfoProps { organizationId: string; @@ -33,7 +35,6 @@ interface OrganizationInfoProps { is_org_admin: boolean; is_proxy_admin: boolean; userModels: string[]; - editOrg: boolean; } const OrganizationInfoView: React.FC = ({ @@ -43,7 +44,6 @@ const OrganizationInfoView: React.FC = ({ is_org_admin, is_proxy_admin, userModels, - editOrg, }) => { const queryClient = useQueryClient(); const { data: orgData, isLoading: loading } = useOrganization(organizationId); @@ -53,10 +53,16 @@ const OrganizationInfoView: React.FC = ({ const [selectedEditMember, setSelectedEditMember] = useState(null); const canEditOrg = is_org_admin || is_proxy_admin; const { data: teams } = useTeams(); - const { onTabChange, hasVisited } = useVisitedTabs(editOrg ? "settings" : "overview"); + const [tab, setTab] = useUrlTab(ORGANIZATION_TABS, "overview", ORGANIZATION_TAB_URL_KEY); + const { onTabChange, hasVisited } = useVisitedTabs(tab); const teamAliasMap = useMemo(() => createTeamAliasMap(teams), [teams]); + const handleTabChange = (value: OrganizationTab) => { + setTab(value); + onTabChange(value); + }; + const handleMemberAdd = async (values: any) => { try { if (accessToken == null) { @@ -158,7 +164,7 @@ const OrganizationInfoView: React.FC = ({
- + Overview From df2dd9b7f25ed30f147087033ad652cf22ff380c Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 09:58:22 -0700 Subject: [PATCH 165/187] feat(ui): keep projects search and project key table state in the URL The projects list search lives in ?project_search= and its pagination now goes through useUrlTableState, keeping the page and page_size keys. The key table inside a project reads keys_search, keys_page and keys_page_size, resets to its first page on a new search, and no longer snaps a deep-linked page while the key fetch is failing. Closing a project drops its keys_ params so they do not leak into the next project --- .../_components/ProjectKeysSection.test.tsx | 111 +++++++++++++++++- .../_components/ProjectKeysSection.tsx | 20 ++-- .../projects/_components/ProjectKeysTable.tsx | 6 +- .../_components/ProjectsPage.test.tsx | 44 ++++++- .../projects/_components/ProjectsPage.tsx | 16 +-- .../_components/ProjectsTable.test.tsx | 4 +- .../projects/_components/ProjectsTable.tsx | 18 ++- .../_components/useProjectsUrlState.ts | 39 ++++++ 8 files changed, 223 insertions(+), 35 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/projects/_components/useProjectsUrlState.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.test.tsx index 0ba1dcab155..470cee475ee 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.test.tsx @@ -1,5 +1,7 @@ -import { describe, it, expect, vi } from "vitest"; -import { renderWithProviders, screen } from "../../../../../tests/test-utils"; +import { describe, it, expect, vi, beforeEach, type Mock } from "vitest"; +import userEvent from "@testing-library/user-event"; +import type { OnUrlUpdateFunction } from "nuqs/adapters/testing"; +import { fireEvent, renderWithProviders, screen, waitFor } from "../../../../../tests/test-utils"; import { ProjectKeysSection } from "./ProjectKeysSection"; const mockUseKeys = vi.fn(); @@ -70,3 +72,108 @@ describe("ProjectKeysSection", () => { ); }); }); + +describe("ProjectKeysSection URL state (keys_ prefix)", () => { + const fortyTwoKeys = { + data: { keys: [], total_count: 42, current_page: 1, total_pages: 9 }, + isLoading: false, + isError: false, + }; + const lastSearchParams = (onUrlUpdate: Mock) => onUrlUpdate.mock.calls.at(-1)?.[0].searchParams; + + beforeEach(() => { + mockUseKeys.mockReset(); + }); + + it("should fetch the page, page size and key name filter named by the keys_ params", () => { + mockUseKeys.mockReturnValue(fortyTwoKeys); + renderWithProviders(, { + searchParams: "?page=4&keys_page=2&keys_page_size=10&keys_search=prod", + }); + + expect(mockUseKeys).toHaveBeenLastCalledWith( + 2, + 10, + expect.objectContaining({ projectID: "proj-1", selectedKeyAlias: "prod" }), + ); + expect(screen.getByPlaceholderText("Filter by key name...")).toHaveValue("prod"); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 2 of 5"); + }); + + it("should write the key name filter to ?keys_search= and return the keys to their first page", async () => { + mockUseKeys.mockReturnValue(fortyTwoKeys); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { + searchParams: "?page=4&keys_page=3", + onUrlUpdate, + }); + + fireEvent.change(screen.getByPlaceholderText("Filter by key name..."), { target: { value: "prod" } }); + + await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("keys_search")).toBe("prod")); + expect(lastSearchParams(onUrlUpdate)?.has("keys_page")).toBe(false); + expect(lastSearchParams(onUrlUpdate)?.get("page")).toBe("4"); + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 5, expect.objectContaining({ selectedKeyAlias: "prod" })); + }); + + it("should remove ?keys_search= when the key filter is cleared", async () => { + const user = userEvent.setup(); + mockUseKeys.mockReturnValue(fortyTwoKeys); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: "?keys_search=prod", onUrlUpdate }); + + await user.click(screen.getByRole("button", { name: /clear key filter/i })); + + await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.has("keys_search")).toBe(false)); + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 5, expect.objectContaining({ selectedKeyAlias: null })); + }); + + it("should write key pages to ?keys_page= without touching the projects list's ?page=", async () => { + const user = userEvent.setup(); + mockUseKeys.mockReturnValue(fortyTwoKeys); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: "?page=4", onUrlUpdate }); + + await user.click(screen.getByTestId("pagination-next")); + + await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("keys_page")).toBe("2")); + expect(lastSearchParams(onUrlUpdate)?.get("page")).toBe("4"); + expect(mockUseKeys).toHaveBeenLastCalledWith(2, 5, expect.anything()); + }); + + it("should snap a ?keys_page= past the last page back to the last page once the keys load", async () => { + mockUseKeys.mockReturnValue({ data: undefined, isLoading: true, isError: false }); + const onUrlUpdate = vi.fn(); + const { rerender } = renderWithProviders(, { + searchParams: "?keys_page=9", + onUrlUpdate, + }); + expect(mockUseKeys).toHaveBeenLastCalledWith(9, 5, expect.anything()); + + mockUseKeys.mockReturnValue({ + data: { keys: [], total_count: 6, current_page: 9, total_pages: 2 }, + isLoading: false, + isError: false, + }); + rerender(); + + await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("keys_page")).toBe("2")); + expect(mockUseKeys).toHaveBeenLastCalledWith(2, 5, expect.anything()); + }); + + it("should keep a deep-linked ?keys_page= when the key fetch fails", async () => { + mockUseKeys.mockReturnValue({ data: undefined, isLoading: true, isError: false }); + const onUrlUpdate = vi.fn(); + const { rerender } = renderWithProviders(, { + searchParams: "?keys_page=3", + onUrlUpdate, + }); + + mockUseKeys.mockReturnValue({ data: undefined, isLoading: false, isError: true }); + rerender(); + + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(onUrlUpdate).not.toHaveBeenCalled(); + expect(mockUseKeys).toHaveBeenLastCalledWith(3, 5, expect.anything()); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.tsx index c618dd6b105..61c8346bce1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.tsx @@ -1,30 +1,27 @@ import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; -import { PaginationState } from "@tanstack/react-table"; import { KeyIcon, SearchIcon, X } from "lucide-react"; -import { useEffect, useState } from "react"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { ProjectKeysTable } from "./ProjectKeysTable"; +import { useProjectKeysTableState } from "./useProjectsUrlState"; interface ProjectKeysSectionProps { projectId: string; } -const PAGE_SIZE = 5; - export function ProjectKeysSection({ projectId }: ProjectKeysSectionProps) { - const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: PAGE_SIZE }); - const [keyAlias, setKeyAlias] = useState(""); + const { + search: keyAlias, + setSearch: setKeyAlias, + pagination, + onPaginationChange: setPagination, + } = useProjectKeysTableState(); - const { data, isLoading } = useKeys(pagination.pageIndex + 1, pagination.pageSize, { + const { data, isLoading, isError } = useKeys(pagination.pageIndex + 1, pagination.pageSize, { projectID: projectId, selectedKeyAlias: keyAlias || null, }); - useEffect(() => { - setPagination((current) => ({ ...current, pageIndex: 0 })); - }, [keyAlias]); - const keys = data?.keys ?? []; const totalCount = data?.total_count ?? 0; @@ -60,6 +57,7 @@ export function ProjectKeysSection({ projectId }: ProjectKeysSectionProps) { keys={keys} totalCount={totalCount} isLoading={isLoading} + isError={isError} pagination={pagination} onPaginationChange={setPagination} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysTable.tsx index 080aad7b26d..53f3898a30f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysTable.tsx @@ -8,16 +8,18 @@ import { KeyResponse } from "@/components/key_team_helpers/key_list"; import { DataTable } from "@/components/shared/DataTable"; import { getProjectKeysTableColumns } from "./ProjectKeysTableColumns"; +import { PROJECT_KEYS_DEFAULT_PAGE_SIZE } from "./useProjectsUrlState"; interface ProjectKeysTableProps { keys: KeyResponse[]; totalCount: number; isLoading: boolean; + isError?: boolean; pagination: PaginationState; onPaginationChange: OnChangeFn; } -const PAGE_SIZE_OPTIONS = [5, 10, 25]; +const PAGE_SIZE_OPTIONS = [PROJECT_KEYS_DEFAULT_PAGE_SIZE, 10, 25]; function EmptyState() { return ( @@ -35,6 +37,7 @@ export function ProjectKeysTable({ keys, totalCount, isLoading, + isError = false, pagination, onPaginationChange, }: ProjectKeysTableProps) { @@ -51,6 +54,7 @@ export function ProjectKeysTable({ rowCount={totalCount} pageSizeOptions={PAGE_SIZE_OPTIONS} isLoading={isLoading} + isError={isError} loadingMessage="Loading keys…" noDataMessage={} size="compact" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.test.tsx index 66d7413f500..2a8f7d11d8f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.test.tsx @@ -190,22 +190,47 @@ describe("ProjectsPage", () => { it("should reset to the first page when the search text changes", async () => { const user = userEvent.setup(); + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); const manyProjects = Array.from({ length: 12 }, (_, i) => ({ ...mockProjects[0], project_id: `proj-${i + 1}`, project_alias: `Project ${String(i + 1).padStart(2, "0")}`, })); mockUseProjects.mockReturnValue({ data: manyProjects, isLoading: false }); - renderWithProviders(); + renderWithProviders(, { onUrlUpdate }); await user.click(screen.getByTestId("pagination-next")); expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 2 of 2"); + await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("page")).toBe("2")); fireEvent.change(screen.getByPlaceholderText(/search projects/i), { target: { value: "Project 01" } }); await waitFor(() => { expect(screen.getByText("Project 01")).toBeInTheDocument(); expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 1"); }); + await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].queryString).toBe("?project_search=Project+01")); + }); + + it("should restore the search box and filtered list from a ?project_search= deep link", () => { + mockUseProjects.mockReturnValue({ data: mockProjects, isLoading: false }); + renderWithProviders(, { searchParams: "?project_search=Beta" }); + + expect(screen.getByPlaceholderText(/search projects/i)).toHaveValue("Beta"); + expect(screen.getByText("Beta Project")).toBeInTheDocument(); + expect(screen.queryByText("Alpha Project")).not.toBeInTheDocument(); + }); + + it("should remove ?project_search= when the search is cleared", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + mockUseProjects.mockReturnValue({ data: mockProjects, isLoading: false }); + renderWithProviders(, { searchParams: "?project_search=Beta", onUrlUpdate }); + + await user.click(screen.getByRole("button", { name: /clear search/i })); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenLastCalledWith(expect.objectContaining({ queryString: "" }))); + expect(screen.getByPlaceholderText(/search projects/i)).toHaveValue(""); + expect(screen.getByText("Alpha Project")).toBeInTheDocument(); }); it("should open the detail view directly from a ?project= deep link", () => { @@ -250,6 +275,23 @@ describe("ProjectsPage", () => { expect(screen.getByText("Alpha Project")).toBeInTheDocument(); }); + it("should drop the project's key table state but keep the list's search and page when the detail view is closed", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + mockUseProjects.mockReturnValue({ data: mockProjects, isLoading: false }); + renderWithProviders(, { + searchParams: "?page=2&project_search=Project&project=proj-1&keys_page=3&keys_page_size=10&keys_search=prod", + onUrlUpdate, + }); + + await user.click(screen.getByRole("button", { name: /back to projects/i })); + + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalledTimes(1)); + const [update] = onUrlUpdate.mock.calls[0]; + expect(update.queryString).toBe("?page=2&project_search=Project"); + expect(update.options.history).toBe("replace"); + }); + it("should resolve team alias from the teams list in the Team column", () => { mockUseTeams.mockReturnValue({ data: [{ team_id: "team-1", team_alias: "Engineering", models: [] }], diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx index 4aba2bb627d..2d3c1acf75e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.tsx @@ -9,6 +9,7 @@ import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from " import { CreateProjectModal } from "./ProjectModals/CreateProjectModal"; import { ProjectDetail } from "./ProjectDetailsPage"; import { ProjectsTable } from "./ProjectsTable"; +import { useClearProjectKeysTableState, useProjectsTableState } from "./useProjectsUrlState"; export function ProjectsPage() { const { data: projects, isLoading } = useProjects(); @@ -18,8 +19,9 @@ export function ProjectsPage() { "project", parseAsString.withOptions({ history: "push" }), ); + const clearProjectKeysTableState = useClearProjectKeysTableState(); + const { search: searchText, setSearch: setSearchText } = useProjectsTableState(); const [isCreateModalVisible, setIsCreateModalVisible] = useState(false); - const [searchText, setSearchText] = useState(""); const teamAliasMap = useMemo(() => { const map = new Map(); @@ -44,13 +46,13 @@ export function ProjectsPage() { }); }, [projects, searchText, teamAliasMap]); + const closeProject = () => { + void setSelectedProjectId(null, { history: "replace" }); + clearProjectKeysTableState(); + }; + if (selectedProjectId) { - return ( - void setSelectedProjectId(null, { history: "replace" })} - /> - ); + return ; } return ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.test.tsx index a1b59f6035c..523932007fa 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.test.tsx @@ -73,7 +73,7 @@ describe("ProjectsTable pagination URL state", () => { expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 11-14 of 14"); }); - it("should push ?page=2 onto history when the next page control is clicked", async () => { + it("should write ?page=2 to the URL when the next page control is clicked", async () => { const user = userEvent.setup(); const onUrlUpdate = vi.fn(); renderTable({ onUrlUpdate }); @@ -83,7 +83,7 @@ describe("ProjectsTable pagination URL state", () => { await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); const [update] = onUrlUpdate.mock.calls[0]; expect(update.searchParams.get("page")).toBe("2"); - expect(update.options.history).toBe("push"); + expect(update.searchParams.has("page_size")).toBe(false); expect(firstDataRow().getByText("Project 11")).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.tsx index 74242f3ed45..1d63958faed 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.tsx @@ -2,13 +2,13 @@ import { SortingState } from "@tanstack/react-table"; import { FolderKanban } from "lucide-react"; -import { parseAsInteger, useQueryStates } from "nuqs"; import { useMemo, useState } from "react"; import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects"; import { DataTable, DataTablePagination } from "@/components/shared/DataTable"; import { getProjectsTableColumns } from "./ProjectsTableColumns"; +import { PROJECTS_DEFAULT_PAGE_SIZE, useProjectsTableState } from "./useProjectsUrlState"; interface ProjectsTableProps { projects: ProjectResponse[]; @@ -19,8 +19,7 @@ interface ProjectsTableProps { isTeamsLoading: boolean; } -const DEFAULT_PAGE_SIZE = 10; -const PAGE_SIZE_OPTIONS = [DEFAULT_PAGE_SIZE, 25, 50]; +const PAGE_SIZE_OPTIONS = [PROJECTS_DEFAULT_PAGE_SIZE, 25, 50]; function EmptyState({ isFiltered }: { isFiltered: boolean }) { return ( @@ -47,11 +46,8 @@ export function ProjectsTable({ isTeamsLoading, }: ProjectsTableProps) { const [sorting, setSorting] = useState([]); - const [{ page, page_size }, setPagination] = useQueryStates( - { page: parseAsInteger.withDefault(1), page_size: parseAsInteger.withDefault(DEFAULT_PAGE_SIZE) }, - { history: "push" }, - ); - const pageSize = PAGE_SIZE_OPTIONS.includes(page_size) ? page_size : DEFAULT_PAGE_SIZE; + const { pagination, onPaginationChange } = useProjectsTableState(); + const pageSize = PAGE_SIZE_OPTIONS.includes(pagination.pageSize) ? pagination.pageSize : PROJECTS_DEFAULT_PAGE_SIZE; const columns = useMemo(() => { const deps = { onProjectClick, teamAliasMap, isTeamsLoading }; @@ -59,7 +55,7 @@ export function ProjectsTable({ }, [onProjectClick, teamAliasMap, isTeamsLoading]); const pageCount = Math.max(Math.ceil(projects.length / pageSize), 1); - const pageIndex = page >= 1 && page <= pageCount ? page - 1 : 0; + const pageIndex = pagination.pageIndex < pageCount ? pagination.pageIndex : 0; return ( void setPagination({ page: nextPageIndex + 1 })} - onPageSizeChange={(nextPageSize) => void setPagination({ page_size: nextPageSize, page: null })} + onPageChange={(nextPageIndex) => onPaginationChange({ pageIndex: nextPageIndex, pageSize })} + onPageSizeChange={(nextPageSize) => onPaginationChange({ pageIndex: 0, pageSize: nextPageSize })} pageSizeOptions={PAGE_SIZE_OPTIONS} isLoading={isLoading} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/useProjectsUrlState.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/useProjectsUrlState.ts new file mode 100644 index 00000000000..db88ad7c3a9 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/useProjectsUrlState.ts @@ -0,0 +1,39 @@ +import { useUrlTableState, type UrlTableState, type UrlTableStateOptions } from "@/components/shared/DataTable"; +import { parseAsString, useQueryStates } from "nuqs"; +import { useCallback } from "react"; + +export const PROJECTS_DEFAULT_PAGE_SIZE = 10; +export const PROJECT_KEYS_DEFAULT_PAGE_SIZE = 5; + +const PROJECT_KEYS_URL_PREFIX = "keys_"; +const TABLE_STATE_URL_KEYS = ["search", "sort_by", "sort_order", "page", "page_size"] as const; + +const PROJECTS_TABLE_STATE_OPTIONS: UrlTableStateOptions = { + sortFields: [], + defaultSort: { id: "created_at", desc: true }, + defaultPageSize: PROJECTS_DEFAULT_PAGE_SIZE, + filterColumns: [], + urlKeys: { search: "project_search" }, +}; + +const PROJECT_KEYS_TABLE_STATE_OPTIONS: UrlTableStateOptions = { + sortFields: [], + defaultSort: { id: "created_at", desc: true }, + defaultPageSize: PROJECT_KEYS_DEFAULT_PAGE_SIZE, + maxPageSize: 25, + filterColumns: [], + keyPrefix: PROJECT_KEYS_URL_PREFIX, +}; + +const PROJECT_KEYS_URL_STATE = Object.fromEntries( + TABLE_STATE_URL_KEYS.map((key) => [`${PROJECT_KEYS_URL_PREFIX}${key}`, parseAsString]), +); + +export const useProjectsTableState = (): UrlTableState => useUrlTableState(PROJECTS_TABLE_STATE_OPTIONS); + +export const useProjectKeysTableState = (): UrlTableState => useUrlTableState(PROJECT_KEYS_TABLE_STATE_OPTIONS); + +export function useClearProjectKeysTableState(): () => void { + const [, setProjectKeysUrlState] = useQueryStates(PROJECT_KEYS_URL_STATE); + return useCallback(() => void setProjectKeysUrlState(null), [setProjectKeysUrlState]); +} From ef8e066c7794bffa25311941a4c8718a53769bfd Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 10:25:27 -0700 Subject: [PATCH 166/187] fix(ui): address review on orgs-projects url state Keep /projects list paging as a pushed history entry, validate the project key table page size against its offered options, and clear the key table params through the table-state setters instead of a copied key list. --- .../_components/OrganizationsPanel.test.tsx | 14 ++++ .../_components/ProjectKeysSection.test.tsx | 28 ++++++++ .../projects/_components/ProjectKeysTable.tsx | 6 +- .../_components/ProjectsPage.test.tsx | 4 +- .../_components/ProjectsTable.test.tsx | 4 +- .../_components/useProjectsUrlState.ts | 68 +++++++++++++++---- 6 files changed, 104 insertions(+), 20 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx index 3e87492778a..bda9f3fbd6c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsPanel.test.tsx @@ -189,6 +189,20 @@ describe("OrganizationsPanel - org detail deep link (?org=)", () => { await expectQueryString("?org=org-plain"); }); + it("closing the org detail keeps the list's search, filter, sort and page in the URL", async () => { + renderPanel({ + searchParams: + "?org_search=Acme&filter_org_id=org-7&sort_by=spend&sort_order=asc&page=2&org=org-x&org_tab=members", + }); + + act(() => mockOrgInfoView.mock.calls.at(-1)?.[0].onClose()); + + await expectQueryString("?org_search=Acme&filter_org_id=org-7&sort_by=spend&sort_order=asc&page=2"); + expect(onUrlUpdate).toHaveBeenCalledTimes(1); + expect(screen.getByPlaceholderText("Search by Organization Name")).toHaveValue("Acme"); + expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "org-7", org_alias: "Acme" }); + }); + it("closing the org detail drops ?org_tab= together with ?org=", async () => { renderPanel({ searchParams: "?org=org-from-url&org_tab=members" }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.test.tsx index 470cee475ee..382c34abcca 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysSection.test.tsx @@ -100,6 +100,34 @@ describe("ProjectKeysSection URL state (keys_ prefix)", () => { expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 2 of 5"); }); + it("should cap an oversized ?keys_page_size= at the largest offered page size", () => { + mockUseKeys.mockReturnValue(fortyTwoKeys); + renderWithProviders(, { searchParams: "?keys_page_size=500" }); + + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 25, expect.anything()); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 2"); + }); + + it("should fall back to the default page size for a ?keys_page_size= outside the offered options", () => { + mockUseKeys.mockReturnValue(fortyTwoKeys); + renderWithProviders(, { searchParams: "?keys_page_size=7" }); + + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 5, expect.anything()); + expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 9"); + }); + + it("should drop an unsupported ?keys_page_size= when the user pages forward", async () => { + const user = userEvent.setup(); + mockUseKeys.mockReturnValue(fortyTwoKeys); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: "?keys_page_size=7", onUrlUpdate }); + + await user.click(screen.getByTestId("pagination-next")); + + await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].queryString).toBe("?keys_page=2")); + expect(mockUseKeys).toHaveBeenLastCalledWith(2, 5, expect.anything()); + }); + it("should write the key name filter to ?keys_search= and return the keys to their first page", async () => { mockUseKeys.mockReturnValue(fortyTwoKeys); const onUrlUpdate = vi.fn(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysTable.tsx index 53f3898a30f..50f40057ec2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectKeysTable.tsx @@ -8,7 +8,7 @@ import { KeyResponse } from "@/components/key_team_helpers/key_list"; import { DataTable } from "@/components/shared/DataTable"; import { getProjectKeysTableColumns } from "./ProjectKeysTableColumns"; -import { PROJECT_KEYS_DEFAULT_PAGE_SIZE } from "./useProjectsUrlState"; +import { PROJECT_KEYS_PAGE_SIZE_OPTIONS } from "./useProjectsUrlState"; interface ProjectKeysTableProps { keys: KeyResponse[]; @@ -19,8 +19,6 @@ interface ProjectKeysTableProps { onPaginationChange: OnChangeFn; } -const PAGE_SIZE_OPTIONS = [PROJECT_KEYS_DEFAULT_PAGE_SIZE, 10, 25]; - function EmptyState() { return (
@@ -52,7 +50,7 @@ export function ProjectKeysTable({ pagination={pagination} onPaginationChange={onPaginationChange} rowCount={totalCount} - pageSizeOptions={PAGE_SIZE_OPTIONS} + pageSizeOptions={PROJECT_KEYS_PAGE_SIZE_OPTIONS} isLoading={isLoading} isError={isError} loadingMessage="Loading keys…" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.test.tsx index 2a8f7d11d8f..309da01b295 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsPage.test.tsx @@ -209,6 +209,7 @@ describe("ProjectsPage", () => { expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 1"); }); await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].queryString).toBe("?project_search=Project+01")); + expect(onUrlUpdate).toHaveBeenCalledTimes(2); }); it("should restore the search box and filtered list from a ?project_search= deep link", () => { @@ -280,7 +281,8 @@ describe("ProjectsPage", () => { const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); mockUseProjects.mockReturnValue({ data: mockProjects, isLoading: false }); renderWithProviders(, { - searchParams: "?page=2&project_search=Project&project=proj-1&keys_page=3&keys_page_size=10&keys_search=prod", + searchParams: + "?page=2&project_search=Project&project=proj-1&keys_page=3&keys_page_size=10&keys_search=prod&keys_sort_by=spend&keys_sort_order=asc", onUrlUpdate, }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.test.tsx index 523932007fa..aaecf98d4ab 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectsTable.test.tsx @@ -73,7 +73,7 @@ describe("ProjectsTable pagination URL state", () => { expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 11-14 of 14"); }); - it("should write ?page=2 to the URL when the next page control is clicked", async () => { + it("should push ?page=2 onto history when the next page control is clicked", async () => { const user = userEvent.setup(); const onUrlUpdate = vi.fn(); renderTable({ onUrlUpdate }); @@ -84,6 +84,7 @@ describe("ProjectsTable pagination URL state", () => { const [update] = onUrlUpdate.mock.calls[0]; expect(update.searchParams.get("page")).toBe("2"); expect(update.searchParams.has("page_size")).toBe(false); + expect(update.options.history).toBe("push"); expect(firstDataRow().getByText("Project 11")).toBeInTheDocument(); }); @@ -147,6 +148,7 @@ describe("ProjectsTable pagination URL state", () => { const lastUpdate = onUrlUpdate.mock.calls.at(-1)?.[0]; expect(lastUpdate.searchParams.get("page")).toBeNull(); expect(lastUpdate.searchParams.get("page_size")).toBe("25"); + expect(lastUpdate.options.history).toBe("push"); }); it("should apply both params from a ?page=2&page_size=25 deep link so the restored view matches", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/useProjectsUrlState.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/useProjectsUrlState.ts index db88ad7c3a9..57c1a8fd167 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/useProjectsUrlState.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/useProjectsUrlState.ts @@ -1,12 +1,11 @@ +import { functionalUpdate, type OnChangeFn, type PaginationState } from "@tanstack/react-table"; import { useUrlTableState, type UrlTableState, type UrlTableStateOptions } from "@/components/shared/DataTable"; -import { parseAsString, useQueryStates } from "nuqs"; -import { useCallback } from "react"; +import { parseAsInteger, useQueryStates } from "nuqs"; +import { useCallback, useMemo } from "react"; export const PROJECTS_DEFAULT_PAGE_SIZE = 10; export const PROJECT_KEYS_DEFAULT_PAGE_SIZE = 5; - -const PROJECT_KEYS_URL_PREFIX = "keys_"; -const TABLE_STATE_URL_KEYS = ["search", "sort_by", "sort_order", "page", "page_size"] as const; +export const PROJECT_KEYS_PAGE_SIZE_OPTIONS = [PROJECT_KEYS_DEFAULT_PAGE_SIZE, 10, 25]; const PROJECTS_TABLE_STATE_OPTIONS: UrlTableStateOptions = { sortFields: [], @@ -16,24 +15,65 @@ const PROJECTS_TABLE_STATE_OPTIONS: UrlTableStateOptions = { urlKeys: { search: "project_search" }, }; +const PROJECTS_PAGE_PARAMS = { + page: parseAsInteger.withDefault(1), + page_size: parseAsInteger.withDefault(PROJECTS_DEFAULT_PAGE_SIZE), +}; + const PROJECT_KEYS_TABLE_STATE_OPTIONS: UrlTableStateOptions = { sortFields: [], defaultSort: { id: "created_at", desc: true }, defaultPageSize: PROJECT_KEYS_DEFAULT_PAGE_SIZE, - maxPageSize: 25, + maxPageSize: Math.max(...PROJECT_KEYS_PAGE_SIZE_OPTIONS), filterColumns: [], - keyPrefix: PROJECT_KEYS_URL_PREFIX, + keyPrefix: "keys_", }; -const PROJECT_KEYS_URL_STATE = Object.fromEntries( - TABLE_STATE_URL_KEYS.map((key) => [`${PROJECT_KEYS_URL_PREFIX}${key}`, parseAsString]), -); +export function useProjectsTableState(): UrlTableState { + const tableState = useUrlTableState(PROJECTS_TABLE_STATE_OPTIONS); + const [, setPageParams] = useQueryStates(PROJECTS_PAGE_PARAMS, { history: "push" }); + const { pagination } = tableState; -export const useProjectsTableState = (): UrlTableState => useUrlTableState(PROJECTS_TABLE_STATE_OPTIONS); + const onPaginationChange = useCallback>( + (updaterOrValue) => { + const next = functionalUpdate(updaterOrValue, pagination); + void setPageParams({ page: next.pageIndex + 1, page_size: next.pageSize }); + }, + [pagination, setPageParams], + ); -export const useProjectKeysTableState = (): UrlTableState => useUrlTableState(PROJECT_KEYS_TABLE_STATE_OPTIONS); + return useMemo(() => ({ ...tableState, onPaginationChange }), [tableState, onPaginationChange]); +} + +export function useProjectKeysTableState(): UrlTableState { + const tableState = useUrlTableState(PROJECT_KEYS_TABLE_STATE_OPTIONS); + const { pagination: urlPagination, onPaginationChange: writePagination } = tableState; + const pageSize = PROJECT_KEYS_PAGE_SIZE_OPTIONS.includes(urlPagination.pageSize) + ? urlPagination.pageSize + : PROJECT_KEYS_DEFAULT_PAGE_SIZE; + + const pagination = useMemo( + () => ({ pageIndex: urlPagination.pageIndex, pageSize }), + [urlPagination.pageIndex, pageSize], + ); + + const onPaginationChange = useCallback>( + (updaterOrValue) => writePagination(functionalUpdate(updaterOrValue, pagination)), + [pagination, writePagination], + ); + + return useMemo( + () => ({ ...tableState, pagination, onPaginationChange }), + [tableState, pagination, onPaginationChange], + ); +} export function useClearProjectKeysTableState(): () => void { - const [, setProjectKeysUrlState] = useQueryStates(PROJECT_KEYS_URL_STATE); - return useCallback(() => void setProjectKeysUrlState(null), [setProjectKeysUrlState]); + const { setSearch, onSortingChange, onColumnFiltersChange, onPaginationChange } = useProjectKeysTableState(); + return useCallback(() => { + setSearch(""); + onSortingChange([]); + onColumnFiltersChange([]); + onPaginationChange({ pageIndex: 0, pageSize: PROJECT_KEYS_DEFAULT_PAGE_SIZE }); + }, [setSearch, onSortingChange, onColumnFiltersChange, onPaginationChange]); } From 249a23b09cd63277f905cac86a8808595c77cff2 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 11:06:55 -0700 Subject: [PATCH 167/187] chore(ui): prune stale eslint suppressions for projects page --- ui/litellm-dashboard/eslint-suppressions.json | 5 ----- 1 file changed, 5 deletions(-) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 773854d29e6..51a7a196d0a 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -972,11 +972,6 @@ "count": 2 } }, - "src/app/(dashboard)/projects/_components/ProjectKeysSection.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx": { "react-hooks/set-state-in-effect": { "count": 2 From 02bccfd89f0bdfdb81c4c28862cc6fc61ed55443 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:13:05 +0000 Subject: [PATCH 168/187] fix(streaming): fill text_tokens when reasoning is counted from stream content Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 3 ++- .../test_streaming_chunk_builder_cursor.py | 26 ++++--------------- 2 files changed, 7 insertions(+), 22 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index a8b1f81702c..99c02cdfc6a 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1100,7 +1100,8 @@ class ChunkProcessor: if reasoning_tokens is not None: if returned_usage.completion_tokens_details is None: returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=reasoning_tokens + reasoning_tokens=reasoning_tokens, + text_tokens=max(0, returned_usage.completion_tokens - reasoning_tokens), ) elif ( returned_usage.completion_tokens_details is not None diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py index f4dbb28533f..e925fd9b4a8 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py @@ -257,14 +257,6 @@ class TestAnthropicCursorBug: @pytest.mark.parametrize("placeholder", [1, 3, 8]) def test_interrupted_reasoning_only_stream_estimates_from_reasoning(self, placeholder: int): - """ - message_start placeholders are not always 1 (live Anthropic streams - have been observed sending 1 and 8 for the same prompt), and a thinking - model cut off before message_delta has streamed only reasoning_content. - The recovered usage, including the completion_tokens_details the cost - calculator bills from, must come from that reasoning rather than from - the placeholder. - """ message_start = _make_chunk( usage=Usage( prompt_tokens=100, @@ -292,13 +284,9 @@ class TestAnthropicCursorBug: ) assert response.usage.total_tokens == response.usage.prompt_tokens + reasoning_tokens details = response.usage.completion_tokens_details - assert (details.text_tokens or 0) + details.reasoning_tokens == response.usage.completion_tokens + assert details.text_tokens + details.reasoning_tokens == response.usage.completion_tokens def test_fallback_counts_reasoning_and_text_together(self): - """ - With no usable provider count, the estimate covers everything the - provider generated: reasoning_content plus visible text, not text alone. - """ reasoning = "First I should check whether the input is sorted. " * 10 text = "The list is already sorted, so no work is needed." chunks = [_make_chunk(reasoning_content=reasoning), _make_chunk(content=text)] @@ -306,16 +294,12 @@ class TestAnthropicCursorBug: response = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "Sort it."}]) text_only = litellm.token_counter(model="claude-sonnet-4-6", text=text, count_response_tokens=True) - reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens - assert reasoning_tokens > 0 - assert response.usage.completion_tokens == text_only + reasoning_tokens + details = response.usage.completion_tokens_details + assert details.reasoning_tokens > 0 + assert response.usage.completion_tokens == text_only + details.reasoning_tokens + assert details.text_tokens == text_only def test_lone_usage_event_with_finish_reason_is_trusted(self): - """ - Guardrails rebuild responses from the chunks yielded to the client, - which excludes the un-yielded message_start. A finished stream then has - exactly one usage event (message_delta) and it must be kept as-is. - """ chunks = [ _make_chunk(content="Yes, "), _make_chunk(content="that works."), From a1ad95dbbd5dc9003598278c610b401f797d6403 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:13:41 -0700 Subject: [PATCH 169/187] fix(gemini): read the minimal thinking floor from the cost map and cover the /v1/messages bridge --- .../llms/openai/chat/gpt_5_transformation.py | 4 +- .../vertex_and_google_ai_studio_gemini.py | 6 ++- ...odel_prices_and_context_window_backup.json | 6 +++ litellm/utils.py | 6 +-- model_prices_and_context_window.json | 6 +++ .../llms/openai/test_gpt5_transformation.py | 16 +++---- ...test_vertex_and_google_ai_studio_gemini.py | 47 +++++++++++++++++-- 7 files changed, 71 insertions(+), 20 deletions(-) diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index b02f953425d..1b93df95341 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -4,9 +4,9 @@ from typing import Final import litellm from litellm.utils import ( - _is_explicitly_disabled_factory, _supports_factory, declared_value_factory, + is_explicitly_disabled_factory, ) from .gpt_transformation import OpenAIGPTConfig @@ -192,7 +192,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): Use this for opt-out checks where unknown models should be allowed through. """ - return _is_explicitly_disabled_factory( + return is_explicitly_disabled_factory( model=cls._model_map_lookup_name(model), custom_llm_provider=None, key=f"supports_{level}_reasoning_effort", diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d719d53e19f..7d616c37ec1 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -79,6 +79,7 @@ from litellm.utils import ( CustomStreamWrapper, ModelResponse, is_base64_encoded, + is_explicitly_disabled_factory, supports_reasoning, ) @@ -110,7 +111,6 @@ else: SUPPORTED_REASONING_EFFORTS: Final = ("minimal", "low", "medium", "high", "none", "disable") -GEMINI_FLASH_MODELS_WITHOUT_MINIMAL_THINKING: Final = ("gemini-3.7-flash", "gemini-3.8-flash") def _unsupported_reasoning_effort(reasoning_effort: str) -> UnsupportedParamsError: @@ -865,7 +865,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def _supports_minimal_thinking_level(model: str) -> bool: lowered: Final = model.lower() is_gemini3flash: Final = "gemini-3" in lowered and "flash" in lowered - return is_gemini3flash and not any(m in lowered for m in GEMINI_FLASH_MODELS_WITHOUT_MINIMAL_THINKING) + return is_gemini3flash and not is_explicitly_disabled_factory( + model=model, custom_llm_provider=None, key="supports_minimal_reasoning_effort" + ) @staticmethod def _map_reasoning_effort_to_thinking_level( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9f91cf82f41..43399d5af53 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25159,6 +25159,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -25216,6 +25217,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27081,6 +27083,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27140,6 +27143,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27555,6 +27559,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27612,6 +27617,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, diff --git a/litellm/utils.py b/litellm/utils.py index 18df5e2abf7..b6f4e85a702 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2675,7 +2675,7 @@ def declared_value_factory(model: str, custom_llm_provider: str | None, key: str """Return a string value the model map declares for *key*, or ``None`` when it says nothing. The string-valued sibling of :func:`_supports_factory` and - :func:`_is_explicitly_disabled_factory`, public where those two are not because it is read + :func:`is_explicitly_disabled_factory`, public like the latter because both are read from the provider configs rather than from this module, sharing their ``get_llm_provider`` -> ``_get_model_info_helper`` chain and their unprefixed-twin fallback (#20885), so a provider-prefixed entry that omits the key still answers @@ -2711,7 +2711,7 @@ def declared_value_factory(model: str, custom_llm_provider: str | None, key: str return None -def _is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, key: str) -> bool: +def is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, key: str) -> bool: """Return True only when the model map explicitly sets *key* to ``False``. This is the opt-out mirror of :func:`_supports_factory`. Where @@ -2830,7 +2830,7 @@ def is_vision_explicitly_disabled(model: str, custom_llm_provider: str | None = The opt-out mirror of :func:`supports_vision`: a missing declaration reads as not disabled, so unknown or newly added models stay eligible for image routing. """ - return _is_explicitly_disabled_factory(model, custom_llm_provider, "supports_vision") + return is_explicitly_disabled_factory(model, custom_llm_provider, "supports_vision") def supports_vision(model: str, custom_llm_provider: str | None = None) -> bool: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9f91cf82f41..43399d5af53 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25159,6 +25159,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -25216,6 +25217,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27081,6 +27083,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27140,6 +27143,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27555,6 +27559,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -27612,6 +27617,7 @@ "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index b538fad71a2..ba51209e0d5 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -8,7 +8,7 @@ from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config from litellm.llms.openai.openai import OpenAIConfig from litellm.utils import ( - _is_explicitly_disabled_factory, + is_explicitly_disabled_factory, peek_reasoning_summary_aliases, strip_reasoning_summary_aliases_from_optional_params, ) @@ -524,19 +524,19 @@ def test_gpt5_minimal_explicitly_disabled_check(gpt5_config: OpenAIGPT5Config): def test_is_explicitly_disabled_factory_minimal(): - """_is_explicitly_disabled_factory returns True only for explicit False entries. + """is_explicitly_disabled_factory returns True only for explicit False entries. Verifies the shared helper used by _is_reasoning_effort_level_explicitly_disabled directly — so future changes to the helper are caught without going through the method wrapper. """ key = "supports_minimal_reasoning_effort" - assert _is_explicitly_disabled_factory("gpt-5.4-mini", None, key) - assert _is_explicitly_disabled_factory("gpt-5.4-nano", None, key) - assert _is_explicitly_disabled_factory("openai/gpt-5.4-mini", None, key) - assert _is_explicitly_disabled_factory("gpt-5.4", None, key) - assert _is_explicitly_disabled_factory("gpt-5.4-pro", None, key) - assert not _is_explicitly_disabled_factory("gpt-5.4-turbo-preview", None, key) + assert is_explicitly_disabled_factory("gpt-5.4-mini", None, key) + assert is_explicitly_disabled_factory("gpt-5.4-nano", None, key) + assert is_explicitly_disabled_factory("openai/gpt-5.4-mini", None, key) + assert is_explicitly_disabled_factory("gpt-5.4", None, key) + assert is_explicitly_disabled_factory("gpt-5.4-pro", None, key) + assert not is_explicitly_disabled_factory("gpt-5.4-turbo-preview", None, key) def test_gpt5_unknown_model_passes_through_minimal(config: OpenAIConfig): diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index a1c31689d09..b36c4e6c205 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -5,11 +5,14 @@ from copy import deepcopy from typing import Final, List, cast from unittest.mock import MagicMock, patch +import httpx import pytest from pydantic import BaseModel import litellm from litellm import ModelResponse, completion +from litellm.llms.anthropic.experimental_pass_through.messages import handler as anthropic_messages_handler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -2683,7 +2686,7 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): [ "gemini-3.7-flash", "vertex_ai/gemini-3.8-flash", - "gemini-3.8-flash-preview", + "gemini/gemini-3.8-flash", ], ) @pytest.mark.parametrize( @@ -2691,7 +2694,7 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): [("minimal", True), ("none", False), ("disable", False)], ) def test_gemini_37_38_flash_floor_minimal_thinking_level( - model, reasoning_effort, include_thoughts + local_model_cost_map, model, reasoning_effort, include_thoughts ): result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( reasoning_effort, model @@ -2717,7 +2720,7 @@ def test_gemini_37_38_flash_floor_minimal_thinking_level( ], ) def test_gemini_flash_minimal_thinking_support( - model, reasoning_effort, expected_level, include_thoughts + local_model_cost_map, model, reasoning_effort, expected_level, include_thoughts ): result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( reasoning_effort, model @@ -2727,7 +2730,7 @@ def test_gemini_flash_minimal_thinking_support( assert result["includeThoughts"] is include_thoughts -def test_gemini_38_flash_feature_flag_uses_low_thinking_level(monkeypatch): +def test_gemini_38_flash_feature_flag_uses_low_thinking_level(local_model_cost_map, monkeypatch): monkeypatch.setattr(litellm, "enable_gemini_default_thinking_level_low", True) thinking_param = {"type": "enabled", "budget_tokens": 1024} @@ -2742,7 +2745,7 @@ def test_gemini_38_flash_feature_flag_uses_low_thinking_level(monkeypatch): assert result_36["thinkingLevel"] == "minimal" -def test_gemini_38_flash_public_reasoning_effort_none_uses_low(): +def test_gemini_38_flash_public_reasoning_effort_none_uses_low(local_model_cost_map): result = VertexGeminiConfig().map_openai_params( non_default_params={"reasoning_effort": "none"}, optional_params={}, @@ -2756,6 +2759,40 @@ def test_gemini_38_flash_public_reasoning_effort_none_uses_low(): } +@pytest.mark.asyncio +async def test_gemini_38_flash_messages_bridge_thinking_disabled_sends_low_thinking_level(local_model_cost_map): + captured: dict[str, dict] = {} + + def upstream(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + request=request, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream)) + + await anthropic_messages_handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="gemini/gemini-3.8-flash", + custom_llm_provider="gemini", + thinking={"type": "disabled"}, + api_key="fake-gemini-key", + client=client, + ) + + assert captured["body"]["generationConfig"]["thinkingConfig"] == { + "thinkingLevel": "low", + "includeThoughts": False, + } + + def test_reasoning_effort_dict_format_gemini_3(): """ Test that reasoning_effort works when passed as dict format from OpenAI Agents SDK. From 8ae1f763394bcd5a74cf3bf76ccd3756399d1958 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 22:36:56 +0000 Subject: [PATCH 170/187] feat(rust): scaffold redis cache crate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 55 ++++++++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/cache-redis/Cargo.toml | 11 ++ litellm-rust/crates/cache-redis/src/cache.rs | 127 ++++++++++++++++++ litellm-rust/crates/cache-redis/src/lib.rs | 3 + .../crates/cache-redis/tests/cache.rs | 6 + 6 files changed, 203 insertions(+) create mode 100644 litellm-rust/crates/cache-redis/Cargo.toml create mode 100644 litellm-rust/crates/cache-redis/src/cache.rs create mode 100644 litellm-rust/crates/cache-redis/src/lib.rs create mode 100644 litellm-rust/crates/cache-redis/tests/cache.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7397742369b..d5b94e261e6 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -70,6 +70,12 @@ dependencies = [ "rustversion", ] +[[package]] +name = "arcstr" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" + [[package]] name = "async-compression" version = "0.4.46" @@ -1915,6 +1921,15 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-redis" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "redis", + "serde_json", +] + [[package]] name = "litellm-core" version = "0.1.0" @@ -2140,6 +2155,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.2" @@ -2656,6 +2681,24 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "redis" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed" +dependencies = [ + "arcstr", + "combine", + "itoa", + "num-bigint", + "percent-encoding", + "ryu", + "sha1_smol", + "socket2 0.6.5", + "url", + "xxhash-rust", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3096,6 +3139,12 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + [[package]] name = "sha2" version = "0.10.9" @@ -4182,6 +4231,12 @@ version = "0.13.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" +[[package]] +name = "xxhash-rust" +version = "0.8.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6" + [[package]] name = "yoke" version = "0.8.3" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 879090870d8..eb3413eb89a 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -23,6 +23,7 @@ pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" +redis = "1.7.0" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml new file mode 100644 index 00000000000..2db2ae33840 --- /dev/null +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "litellm-cache-redis" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +redis.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs new file mode 100644 index 00000000000..963783d31a6 --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -0,0 +1,127 @@ +use std::sync::Mutex; +use std::time::Duration; + +use litellm_cache::{ + BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs, + Error, +}; +use redis::Commands; + +const DEFAULT_TTL: Duration = Duration::from_secs(600); + +pub struct RedisCache { + connection: Mutex, + default_ttl: Duration, +} + +impl RedisCache { + pub fn new(url: &str, default_ttl: Option) -> Result { + let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; + let connection = client.get_connection().map_err(|_| Error::Unavailable)?; + Ok(Self { + connection: Mutex::new(connection), + default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), + }) + } + + fn connection(&self) -> Result, Error> { + self.connection.lock().map_err(|_| Error::Unavailable) + } + + fn encode(value: &CacheEntry) -> Result, Error> { + serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) + } + + fn decode(value: Vec) -> Result { + serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry) + } + + fn ttl_seconds(ttl: Duration) -> u64 { + ttl.as_secs().max(1) + } +} + +impl BaseCache for RedisCache { + type Value = CacheEntry; + + fn default_ttl(&self) -> Duration { + self.default_ttl + } + + fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> { + let payload = Self::encode(&value)?; + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + self.connection()? + .set_ex::<_, _, ()>(key, payload, ttl) + .map_err(|_| Error::Unavailable) + } + + fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { + self.connection()? + .get::<_, Option>>(key) + .map_err(|_| Error::Unavailable)? + .map(Self::decode) + .transpose() + } + + fn delete_cache(&self, key: &str) -> Result<(), Error> { + self.connection()? + .del::<_, ()>(key) + .map_err(|_| Error::Unavailable) + } + + fn flush_cache(&self) -> Result<(), Error> { + self.connection()? + .flushdb::<()>() + .map_err(|_| Error::Unavailable) + } + + fn disconnect(&self) -> CacheFuture<'_, ()> { + Box::pin(async { Ok(()) }) + } + + fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> { + Box::pin(async { + let mut connection = self.connection()?; + redis::cmd("PING") + .query::(&mut *connection) + .map_err(|_| Error::Unavailable)?; + Ok(CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, + }) + }) + } +} + +#[cfg(test)] +mod tests { + use super::RedisCache; + use litellm_cache::CacheEntry; + use serde_json::json; + use std::time::Duration; + + #[test] + fn cache_entries_round_trip_through_json() { + let entry = CacheEntry { + timestamp: 123.0, + response: json!({"choices": [{"text": "cached"}]}), + }; + + let encoded = RedisCache::encode(&entry).unwrap(); + assert_eq!(RedisCache::decode(encoded).unwrap(), entry); + } + + #[test] + fn invalid_json_is_rejected() { + assert!(RedisCache::decode(b"not json".to_vec()).is_err()); + } + + #[test] + fn ttl_seconds_keeps_redis_expiration_positive() { + assert_eq!(RedisCache::ttl_seconds(Duration::ZERO), 1); + assert_eq!(RedisCache::ttl_seconds(Duration::from_millis(1500)), 1); + assert_eq!(RedisCache::ttl_seconds(Duration::from_secs(15)), 15); + } +} diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs new file mode 100644 index 00000000000..37b35c5ea4a --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -0,0 +1,3 @@ +mod cache; + +pub use cache::RedisCache; diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs new file mode 100644 index 00000000000..76f73145da8 --- /dev/null +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -0,0 +1,6 @@ +use litellm_cache_redis::RedisCache; + +#[test] +fn constructor_rejects_invalid_urls() { + assert!(RedisCache::new("not a redis url", None).is_err()); +} From 0b3c3885bced09b4e98aaed57b58fca742a20339 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 22:38:38 +0000 Subject: [PATCH 171/187] fix(rust): scope Redis dependency to cache crate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.toml | 1 - litellm-rust/crates/cache-redis/Cargo.toml | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index eb3413eb89a..879090870d8 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -23,7 +23,6 @@ pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" -redis = "1.7.0" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index 2db2ae33840..d954084168f 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,5 +7,5 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis.workspace = true +redis = "1.7.0" serde_json.workspace = true From 2a9fa48730bf4cfe19b5a36d5144d11944b42a5b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:18:44 -0700 Subject: [PATCH 172/187] fix: expand wildcard deployments for proxy admins on /model_group/info --- litellm/proxy/proxy_server.py | 10 +++++++++- .../proxy/proxy_server/test_routes_model_info.py | 16 ++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6fe2ab57112..2c5c635b14b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15542,7 +15542,15 @@ async def model_group_info( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ) all_models_str: Final = ( - llm_router.get_model_names() + get_complete_model_list( + key_models=(), + team_models=(), + proxy_model_list=llm_router.get_model_names(), + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + return_wildcard_routes=False, + llm_router=llm_router, + ) if is_proxy_admin else await get_available_models_for_user( user_api_key_dict=user_api_key_dict, diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 2c101156c6c..3000a2ea101 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -353,6 +353,22 @@ def test_model_group_info_proxy_admin_ignores_key_model_restriction( assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4", "claude-3"] +@pytest.mark.parametrize("admin_role", ["proxy_admin", "proxy_admin_viewer"]) +def test_model_group_info_proxy_admin_expands_wildcard_deployments(client, auth_as, model_group_info_router, admin_role): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + model_group_info_router.get_model_names.return_value = ["gpt-4", "anthropic/*"] + known_anthropic_models = get_known_models_from_wildcard(wildcard_model="anthropic/*") + assert known_anthropic_models + + with auth_as(LitellmUserRoles(admin_role), models=["no-default-models"]): + response = client.get("/model_group/info") + + assert response.status_code == 200 + assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4", *known_anthropic_models] + + def test_model_group_info_internal_user_key_model_restriction_applies(client, auth_as, model_group_info_router): from litellm.proxy._types import LitellmUserRoles From 93ba409adf9215ca051372714add07a5c406f89e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 23:19:28 +0000 Subject: [PATCH 173/187] fix(rust): address Redis cache review findings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 46 ++++ litellm-rust/crates/cache-redis/Cargo.toml | 4 + litellm-rust/crates/cache-redis/src/cache.rs | 252 ++++++++++++++++--- 3 files changed, 270 insertions(+), 32 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index d5b94e261e6..9cfce7e0704 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1843,6 +1843,12 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litellm-auth" version = "0.1.0" @@ -1927,7 +1933,9 @@ version = "0.1.0" dependencies = [ "litellm-cache", "redis", + "redis-test", "serde_json", + "tokio", ] [[package]] @@ -2699,6 +2707,18 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "redis-test" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca" +dependencies = [ + "rand 0.9.5", + "redis", + "socket2 0.6.5", + "tempfile", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2889,6 +2909,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.21.12" @@ -3348,6 +3381,19 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" version = "1.0.69" diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index d954084168f..933b0feaae4 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -9,3 +9,7 @@ repository.workspace = true litellm-cache.workspace = true redis = "1.7.0" serde_json.workspace = true +tokio.workspace = true + +[dev-dependencies] +redis-test = "1.0.4" diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 963783d31a6..69dee6c6363 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -1,4 +1,4 @@ -use std::sync::Mutex; +use std::sync::{Arc, Mutex, MutexGuard}; use std::time::Duration; use litellm_cache::{ @@ -8,26 +8,45 @@ use litellm_cache::{ use redis::Commands; const DEFAULT_TTL: Duration = Duration::from_secs(600); +const KEY_PREFIX: &str = "litellm-cache:"; -pub struct RedisCache { - connection: Mutex, +pub struct RedisCache { + connection: Arc>, default_ttl: Duration, } -impl RedisCache { +impl RedisCache { pub fn new(url: &str, default_ttl: Option) -> Result { let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; let connection = client.get_connection().map_err(|_| Error::Unavailable)?; - Ok(Self { - connection: Mutex::new(connection), + Ok(Self::with_connection(connection, default_ttl)) + } +} + +impl RedisCache +where + C: redis::ConnectionLike + Send + 'static, +{ + fn with_connection(connection: C, default_ttl: Option) -> Self { + Self { + connection: Arc::new(Mutex::new(connection)), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), - }) + } } - fn connection(&self) -> Result, Error> { + fn connection(&self) -> Result, Error> { self.connection.lock().map_err(|_| Error::Unavailable) } + fn namespaced_key(key: &str) -> String { + format!("{KEY_PREFIX}{key}") + } + + fn namespaced_pattern() -> &'static str { + const PATTERN: &str = "litellm-cache:*"; + PATTERN + } + fn encode(value: &CacheEntry) -> Result, Error> { serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) } @@ -37,11 +56,31 @@ impl RedisCache { } fn ttl_seconds(ttl: Duration) -> u64 { - ttl.as_secs().max(1) + ttl.as_secs() + .saturating_add(u64::from(ttl.subsec_nanos() > 0)) + .max(1) + } + + fn run_blocking(connection: Arc>, operation: F) -> CacheFuture<'static, T> + where + T: Send + 'static, + F: FnOnce(&mut C) -> Result + Send + 'static, + { + Box::pin(async move { + tokio::task::spawn_blocking(move || { + let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + operation(&mut connection) + }) + .await + .map_err(|_| Error::Unavailable)? + }) } } -impl BaseCache for RedisCache { +impl BaseCache for RedisCache +where + C: redis::ConnectionLike + Send + 'static, +{ type Value = CacheEntry; fn default_ttl(&self) -> Duration { @@ -52,13 +91,13 @@ impl BaseCache for RedisCache { let payload = Self::encode(&value)?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); self.connection()? - .set_ex::<_, _, ()>(key, payload, ttl) + .set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl) .map_err(|_| Error::Unavailable) } fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { self.connection()? - .get::<_, Option>>(key) + .get::<_, Option>>(Self::namespaced_key(key)) .map_err(|_| Error::Unavailable)? .map(Self::decode) .transpose() @@ -66,26 +105,101 @@ impl BaseCache for RedisCache { fn delete_cache(&self, key: &str) -> Result<(), Error> { self.connection()? - .del::<_, ()>(key) + .del::<_, ()>(Self::namespaced_key(key)) .map_err(|_| Error::Unavailable) } fn flush_cache(&self) -> Result<(), Error> { - self.connection()? - .flushdb::<()>() + let mut connection = self.connection()?; + let keys = connection + .scan_match(Self::namespaced_pattern()) + .map_err(|_| Error::Unavailable)? + .collect::>>() + .map_err(|_| Error::Unavailable)?; + if keys.is_empty() { + return Ok(()); + } + connection + .del::<_, usize>(keys) + .map(|_| ()) .map_err(|_| Error::Unavailable) } + fn async_set_cache<'a>( + &'a self, + key: &'a str, + value: Self::Value, + kwargs: CacheKwargs, + ) -> CacheFuture<'a, ()> { + let payload = Self::encode(&value); + let key = Self::namespaced_key(key); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + Self::run_blocking(Arc::clone(&self.connection), move |connection| { + connection + .set_ex::<_, _, ()>(key, payload?, ttl) + .map_err(|_| Error::Unavailable) + }) + } + + fn async_get_cache<'a>( + &'a self, + key: &'a str, + _: &'a CacheKwargs, + ) -> CacheFuture<'a, Option> { + let key = Self::namespaced_key(key); + Box::pin(async move { + Self::run_blocking(Arc::clone(&self.connection), move |connection| { + connection + .get::<_, Option>>(key) + .map_err(|_| Error::Unavailable) + }) + .await? + .map(Self::decode) + .transpose() + }) + } + + fn async_set_cache_pipeline<'a>( + &'a self, + cache_list: Vec<(String, Self::Value)>, + kwargs: CacheKwargs, + ) -> CacheFuture<'a, ()> { + let entries = cache_list + .into_iter() + .map(|(key, value)| { + Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload)) + }) + .collect::, _>>(); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + Self::run_blocking(Arc::clone(&self.connection), move |connection| { + for (key, payload) in entries? { + connection + .set_ex::<_, _, ()>(key, payload, ttl) + .map_err(|_| Error::Unavailable)?; + } + Ok(()) + }) + } + + fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> { + let key = Self::namespaced_key(key); + Self::run_blocking(Arc::clone(&self.connection), move |connection| { + connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) + }) + } + fn disconnect(&self) -> CacheFuture<'_, ()> { Box::pin(async { Ok(()) }) } fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> { - Box::pin(async { - let mut connection = self.connection()?; - redis::cmd("PING") - .query::(&mut *connection) - .map_err(|_| Error::Unavailable)?; + Box::pin(async move { + Self::run_blocking(Arc::clone(&self.connection), |connection| { + redis::cmd("PING") + .query::(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; Ok(CacheConnectionResult { status: CacheConnectionStatus::Success, message: "Redis cache connection test successful".into(), @@ -98,30 +212,104 @@ impl BaseCache for RedisCache { #[cfg(test)] mod tests { use super::RedisCache; - use litellm_cache::CacheEntry; + use litellm_cache::{BaseCache, CacheEntry, CacheKwargs}; + use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; use std::time::Duration; - #[test] - fn cache_entries_round_trip_through_json() { - let entry = CacheEntry { + fn entry() -> CacheEntry { + CacheEntry { timestamp: 123.0, response: json!({"choices": [{"text": "cached"}]}), - }; + } + } - let encoded = RedisCache::encode(&entry).unwrap(); - assert_eq!(RedisCache::decode(encoded).unwrap(), entry); + #[test] + fn cache_entries_round_trip_through_json() { + let entry = entry(); + let encoded = RedisCache::::encode(&entry).unwrap(); + assert_eq!( + RedisCache::::decode(encoded).unwrap(), + entry + ); } #[test] fn invalid_json_is_rejected() { - assert!(RedisCache::decode(b"not json".to_vec()).is_err()); + assert!(RedisCache::::decode(b"not json".to_vec()).is_err()); } #[test] - fn ttl_seconds_keeps_redis_expiration_positive() { - assert_eq!(RedisCache::ttl_seconds(Duration::ZERO), 1); - assert_eq!(RedisCache::ttl_seconds(Duration::from_millis(1500)), 1); - assert_eq!(RedisCache::ttl_seconds(Duration::from_secs(15)), 15); + fn ttl_seconds_rounds_up_and_keeps_expiration_positive() { + assert_eq!( + RedisCache::::ttl_seconds(Duration::ZERO), + 1 + ); + assert_eq!( + RedisCache::::ttl_seconds(Duration::from_millis(1500)), + 2 + ); + assert_eq!( + RedisCache::::ttl_seconds(Duration::from_secs(15)), + 15 + ); + } + + #[test] + fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() { + let value = entry(); + let payload = RedisCache::::encode(&value).unwrap(); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SETEX") + .arg("litellm-cache:key") + .arg(600) + .arg(payload.clone()), + Ok("OK"), + ), + MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)), + MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None); + + cache + .set_cache("key", value.clone(), CacheKwargs::default()) + .unwrap(); + assert_eq!( + cache.get_cache("key", &CacheKwargs::default()).unwrap(), + Some(value) + ); + cache.delete_cache("key").unwrap(); + } + + #[test] + fn flush_scans_and_deletes_only_cache_keys() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(0) + .arg("MATCH") + .arg("litellm-cache:*"), + Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])), + ), + MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None); + + cache.flush_cache().unwrap(); + } + + #[tokio::test] + async fn test_connection_runs_ping_off_executor() { + let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None); + + assert_eq!( + cache.test_connection().await.unwrap().status, + litellm_cache::CacheConnectionStatus::Success + ); } } From ab92a6637d1308af6215edda3a50582548d639bf Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 16:26:52 -0700 Subject: [PATCH 174/187] test(e2e): cover team admin editable fields on /team/update Team admins are refused until a proxy admin enables a field, then limited to the enabled fields, and resending unchanged budget settings keeps the team's budget reset times --- tests/e2e/coverage_registry/mgmt.yaml | 3 + .../management/test_team_management_e2e.py | 299 +++++++++++++++++- 2 files changed, 292 insertions(+), 10 deletions(-) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 31ad61ba3e2..d93d2b2cc67 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -30,6 +30,9 @@ - {id: mgmt.key.health.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4292", rationale: "Key health endpoint"} - {id: mgmt.key.bulk_update.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:2677", rationale: "Batch key updates"} - {id: mgmt.team.update.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1582", rationale: "Metadata/budget updates persist"} +- {id: mgmt.team.update.team_admin_forbidden_until_enabled, module: mgmt, tier: P0, surface: api, assertions: [team_admin_forbidden_until_enabled], source: "team_admin_field_permissions.py:156", rationale: "With no team admin editable fields enabled, a team admin's /team/update is 403 and /team/info reports editing disabled"} +- {id: mgmt.team.update.team_admin_limited_to_enabled_fields, module: mgmt, tier: P0, surface: api, assertions: [team_admin_limited_to_enabled_fields], source: "team_admin_field_permissions.py:156", rationale: "A team admin may change only the enabled fields; a request that also changes any other field is 403 and writes nothing"} +- {id: mgmt.team.update.team_admin_resend_keeps_budget_reset, module: mgmt, tier: P1, surface: api, assertions: [team_admin_resend_keeps_budget_reset], source: "team_admin_field_permissions.py:147", fail_before_fix: proven, rationale: "A team admin resending unchanged budget settings with an enabled field must not push the team's budget reset times back"} - {id: mgmt.team.delete.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1750", rationale: "Deletion prevents key access"} - {id: mgmt.team.block.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py", rationale: "Block suspends all members"} - {id: mgmt.team.info.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "team_endpoints.py:2244", rationale: "Metadata+members+budgets"} diff --git a/tests/e2e/management/test_team_management_e2e.py b/tests/e2e/management/test_team_management_e2e.py index 108aeaad21b..f21931b6ff1 100644 --- a/tests/e2e/management/test_team_management_e2e.py +++ b/tests/e2e/management/test_team_management_e2e.py @@ -1,5 +1,6 @@ """Live e2e: the /team/* management routes' block, membership, and admin-only -contract. +contract, plus the team settings a team admin may change on /team/update once a +proxy admin enables them under Settings > UI > Team admin editable fields. Each test creates its team/user/key resources under unique names (deleted on teardown) and asserts both halves of the contract: the recorded state (the info @@ -8,21 +9,25 @@ Team writes reach the read path once their db/cache entry propagates, so the read-backs poll to a deadline instead of asserting once. Everything the shared harness does not already model lives here: the local -request/response models for /team/block, /team/member_update, and the -/team/info fields (blocked flag and per-member budget) these tests assert on. +request/response models for /team/block, /team/member_update, the partial +/team/update, the UI settings allow-list, and the /team/info fields (blocked +flag, limits, budgets, per-member budget, the caller's edit access) these tests +assert on. """ from __future__ import annotations import time -from collections.abc import Callable -from typing import Literal +from collections.abc import Callable, Generator +from contextlib import contextmanager +from datetime import UTC, datetime, timedelta +from typing import Final, Literal import pytest from pydantic import BaseModel -from e2e_config import unique_marker -from e2e_http import NoBody, StreamingResponse, unwrap +from e2e_config import settle_propagation, unique_marker +from e2e_http import NoBody, PartialBody, StreamingResponse, unwrap from lifecycle import ResourceManager from management_client import ManagementClient from models import ( @@ -39,6 +44,8 @@ pytestmark = pytest.mark.e2e TeamRole = Literal["admin", "user"] +_TEAM_TPM_LIMIT: Final = 1000 + class TeamBlockBody(BaseModel): team_id: str @@ -66,11 +73,37 @@ class TeamMembership(BaseModel): litellm_budget_table: MemberBudgetTable | None = None -class TeamInfoData(BaseModel): +class CallerEditAccess(BaseModel): + kind: Literal["unrestricted", "team_admin", "team_admin_disabled", "none"] + editable_fields: list[str] = [] + + +class BudgetWindow(BaseModel): + budget_duration: str + max_budget: float + reset_at: str | None = None + + +class TeamCustomMetadata(BaseModel): + cost_center: str | None = None + + +class TeamSettings(BaseModel): team_alias: str | None = None models: list[str] = [] + tpm_limit: int | None = None + rpm_limit: int | None = None + max_budget: float | None = None + budget_duration: str | None = None + budget_limits: list[BudgetWindow] | None = None + metadata: TeamCustomMetadata | None = None + + +class TeamInfoData(TeamSettings): blocked: bool | None = None members_with_roles: list[MemberRoleEntry] = [] + budget_reset_at: datetime | None = None + caller_edit_access: CallerEditAccess | None = None class TeamInfoRead(BaseModel): @@ -79,6 +112,27 @@ class TeamInfoRead(BaseModel): team_memberships: list[TeamMembership] = [] +class TeamWithAdminNewBody(TeamNewBody): + tpm_limit: int + members_with_roles: list[TeamMemberEntry] + + +class TeamSettingsChange(PartialBody, TeamSettings): + pass + + +class TeamSettingsUpdate(TeamSettingsChange): + team_id: str + + +class TeamAdminEditableFields(BaseModel): + team_admin_editable_team_fields: list[str] = [] + + +class UiSettingsRead(BaseModel): + values: TeamAdminEditableFields + + def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T: deadline = time.monotonic() + client.proxy.poll_timeout while time.monotonic() < deadline: @@ -107,17 +161,27 @@ def _generate_key(client: ManagementClient, resources: ResourceManager, body: Ke return key -def _read_team(client: ManagementClient, team_id: str) -> TeamInfoRead: +def _read_team(client: ManagementClient, team_id: str, caller_key: str | None = None) -> TeamInfoRead: return unwrap( client.proxy.transport.get( "/team/info", - headers=client.proxy.transport.master, + headers=client.proxy.transport.master if caller_key is None else client.proxy.transport.bearer(caller_key), params=TeamInfoParams(team_id=team_id), response_type=TeamInfoRead, ) ) +def _poll_team( + client: ManagementClient, team_id: str, ready: Callable[[TeamInfoData], bool], failure: str +) -> TeamInfoData: + def read() -> TeamInfoData | None: + info = _read_team(client, team_id).team_info + return info if ready(info) else None + + return _poll(client, read, failure) + + def _set_blocked(client: ManagementClient, team_id: str, *, blocked: bool) -> None: _ = unwrap( client.proxy.transport.post( @@ -301,3 +365,218 @@ class TestTeamManagementRoutes: client.add_team_member(team_id, member_id) member_key = _generate_key(client, resources, KeyGenerateBody(user_id=member_id, team_id=team_id)) return member_id, other_id, member_key, team_id + + +def _team_admin_editable_fields(client: ManagementClient) -> list[str]: + return unwrap( + client.proxy.transport.get( + "/get/ui_settings", + headers=client.proxy.transport.master, + params=NoBody(), + response_type=UiSettingsRead, + ) + ).values.team_admin_editable_team_fields + + +def _set_team_admin_editable_fields(client: ManagementClient, fields: list[str]) -> None: + _ = unwrap( + client.proxy.transport.patch( + "/update/ui_settings", + headers=client.proxy.transport.master, + json=TeamAdminEditableFields(team_admin_editable_team_fields=fields), + response_type=NoBody, + ) + ) + + +@contextmanager +def _team_admins_may_edit(client: ManagementClient, fields: list[str]) -> Generator[None]: + """The allow-list is proxy-wide, so restore whatever was there. Other replicas pick a change up on their + config reload, which the wait covers before any team admin call lands on one of them.""" + original = _team_admin_editable_fields(client) + _set_team_admin_editable_fields(client, fields) + settle_propagation(time.monotonic()) + try: + yield + finally: + _set_team_admin_editable_fields(client, original) + + +@pytest.fixture(scope="class") +def no_team_admin_editable_fields(client: ManagementClient) -> Generator[None]: + with _team_admins_may_edit(client, []): + yield + + +@pytest.fixture(scope="class") +def tpm_limit_editable_by_team_admins(client: ManagementClient) -> Generator[None]: + with _team_admins_may_edit(client, ["tpm_limit"]): + yield + + +def _team_with_admin(client: ManagementClient, resources: ResourceManager) -> tuple[str, str]: + """A team with a tpm_limit, and the key of a user who is an admin of that team.""" + admin_id = _create_user(client, resources, f"e2e-team-admin-{unique_marker()}@example.com") + team_id = client.create_team( + TeamWithAdminNewBody( + team_alias=f"e2e-team-admin-{unique_marker()}", + tpm_limit=_TEAM_TPM_LIMIT, + members_with_roles=[TeamMemberEntry(role="admin", user_id=admin_id)], + ) + ) + resources.defer(lambda: client.delete_team(team_id)) + return team_id, _generate_key(client, resources, KeyGenerateBody(user_id=admin_id)) + + +def _update_team_as(client: ManagementClient, caller_key: str, body: TeamSettingsUpdate) -> StreamingResponse: + return client.proxy.transport.send("/team/update", headers=client.proxy.transport.bearer(caller_key), json=body) + + +@pytest.mark.usefixtures("no_team_admin_editable_fields") +class TestTeamAdminWithNoEditableFields: + """No proxy admin has enabled a team field for team admins, which is how every proxy starts.""" + + @pytest.mark.covers("mgmt.team.update.team_admin_forbidden_until_enabled") + def test_team_admin_cannot_change_any_team_setting( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + team_id, admin_key = _team_with_admin(client, resources) + access = _read_team(client, team_id, admin_key).team_info.caller_edit_access + assert access == CallerEditAccess(kind="team_admin_disabled"), ( + f"/team/info should tell the team admin that editing is disabled, got {access}" + ) + + outcome = _update_team_as(client, admin_key, TeamSettingsUpdate(team_id=team_id, tpm_limit=5000)) + + assert outcome.status_code == 403, ( + f"/team/update by a team admin must be 403 while nothing is enabled, got {outcome.status_code}: " + f"{outcome.body[:300]}" + ) + assert "cannot edit team settings" in outcome.body, f"403 body should say why, got: {outcome.body[:300]}" + tpm_limit = _read_team(client, team_id).team_info.tpm_limit + assert tpm_limit == _TEAM_TPM_LIMIT, f"the refused update still changed tpm_limit to {tpm_limit}" + + +@pytest.mark.usefixtures("tpm_limit_editable_by_team_admins") +class TestTeamAdminWithTpmLimitEnabled: + """A proxy admin has enabled tpm_limit, so a team admin may change that setting and no other.""" + + @pytest.mark.covers("mgmt.team.update.team_admin_limited_to_enabled_fields") + def test_team_admin_saves_the_settings_form_with_a_new_tpm_limit( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + team_id, admin_key = _team_with_admin(client, resources) + access = _read_team(client, team_id, admin_key).team_info.caller_edit_access + assert access == CallerEditAccess(kind="team_admin", editable_fields=["tpm_limit"]), ( + f"/team/info should list tpm_limit as the team admin's only editable field, got {access}" + ) + before = _read_team(client, team_id).team_info + + outcome = _update_team_as( + client, + admin_key, + TeamSettingsUpdate(team_id=team_id, team_alias=before.team_alias, models=before.models, tpm_limit=5000), + ) + + assert outcome.status_code == 200, ( + f"a team admin resending the form with only tpm_limit changed must succeed, got {outcome.status_code}: " + f"{outcome.body[:300]}" + ) + after = _poll_team( + client, team_id, lambda info: info.tpm_limit == 5000, "/team/info never reflected tpm_limit=5000" + ) + assert after.model_copy(update={"tpm_limit": _TEAM_TPM_LIMIT}) == before, ( + f"the update changed more than tpm_limit: before {before}, after {after}" + ) + + @pytest.mark.covers("mgmt.team.update.team_admin_limited_to_enabled_fields") + @pytest.mark.parametrize( + "change", + [ + pytest.param(TeamSettingsChange(rpm_limit=10), id="rpm_limit"), + pytest.param(TeamSettingsChange(max_budget=0.5), id="max_budget"), + pytest.param(TeamSettingsChange(team_alias="renamed-by-team-admin"), id="team_alias"), + pytest.param(TeamSettingsChange(models=["gemini-2.5-flash"]), id="models"), + pytest.param(TeamSettingsChange(budget_duration="1d"), id="budget_duration"), + pytest.param(TeamSettingsChange(metadata=TeamCustomMetadata(cost_center="team-admin")), id="metadata"), + ], + ) + def test_team_admin_cannot_change_a_setting_that_is_not_enabled( + self, client: ManagementClient, resources: ResourceManager, change: TeamSettingsChange + ) -> None: + (field,) = change.model_fields_set + team_id, admin_key = _team_with_admin(client, resources) + before = _read_team(client, team_id).team_info + + outcome = _update_team_as( + client, + admin_key, + TeamSettingsUpdate.model_validate( + {**change.model_dump(exclude_unset=True), "team_id": team_id, "tpm_limit": 5000} + ), + ) + + assert outcome.status_code == 403, ( + f"a team admin changing {field} must be 403, got {outcome.status_code}: {outcome.body[:300]}" + ) + assert f"'{field}'" in outcome.body, f"403 body should name {field}, got: {outcome.body[:300]}" + after = _read_team(client, team_id).team_info + assert after == before, ( + f"the refused update still wrote to the team, the enabled tpm_limit included: before {before}, " + f"after {after}" + ) + + @pytest.mark.covers("mgmt.team.update.team_admin_resend_keeps_budget_reset") + def test_team_admin_resending_the_budget_settings_keeps_the_next_budget_reset( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + """A 120s budget resets at the start of the minute after next. Resending it once the next minute has + started would push that reset a minute later, while the stored reset is still a minute out, so the + proxy's budget reset job cannot be what moves it.""" + team_id, admin_key = _team_with_admin(client, resources) + _ = unwrap( + client.proxy.transport.post( + "/team/update", + headers=client.proxy.transport.master, + json=TeamSettingsUpdate( + team_id=team_id, + budget_duration="120s", + budget_limits=[BudgetWindow(budget_duration="120s", max_budget=5.0)], + ), + response_type=NoBody, + ) + ) + budgeted = _poll_team( + client, + team_id, + lambda info: info.budget_reset_at is not None and bool(info.budget_limits), + "/team/info never reflected the 120s budget the proxy admin set", + ) + assert budgeted.budget_reset_at is not None + next_minute = budgeted.budget_reset_at - timedelta(seconds=58) + time.sleep(max(0.0, (next_minute - datetime.now(UTC)).total_seconds())) + + outcome = _update_team_as( + client, + admin_key, + TeamSettingsUpdate( + team_id=team_id, + tpm_limit=5000, + budget_duration=budgeted.budget_duration, + budget_limits=budgeted.budget_limits, + ), + ) + + assert outcome.status_code == 200, ( + f"resending unchanged budget settings with a new tpm_limit must succeed, got {outcome.status_code}: " + f"{outcome.body[:300]}" + ) + after = _poll_team( + client, team_id, lambda info: info.tpm_limit == 5000, "/team/info never reflected tpm_limit=5000" + ) + assert after.budget_reset_at == budgeted.budget_reset_at, ( + f"the team admin pushed the budget reset from {budgeted.budget_reset_at} to {after.budget_reset_at}" + ) + assert after.budget_limits == budgeted.budget_limits, ( + f"the team admin pushed the budget window resets from {budgeted.budget_limits} to {after.budget_limits}" + ) From 15c18ad6cdd6c9cc70e6ad1f5e31c3e5c36c4da5 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:27:07 +0000 Subject: [PATCH 175/187] fix(bedrock_mantle): accept verbosity on gpt-5.x chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock_mantle/chat/transformation.py | 3 +++ .../test_bedrock_mantle_transformation.py | 12 ++++++++++++ 2 files changed, 15 insertions(+) diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index a1153dffc93..5b69d7aff42 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -22,6 +22,7 @@ from litellm.llms.bedrock_mantle.common_utils import ( BEDROCK_MANTLE_DEFAULT_REGION, BedrockMantleAuthMixin, ) +from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams @@ -108,6 +109,8 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): def get_supported_openai_params(self, model: str) -> list: base_params: Final = super().get_supported_openai_params(model) + if is_gpt_reasoning_series_name(model) and "verbosity" not in base_params: + base_params.append("verbosity") try: if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): if "reasoning_effort" not in base_params: diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index c948dfb3553..15570eaec4d 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -257,6 +257,18 @@ class TestBedrockMantleConfig: assert "temperature" in params assert "stream" in params assert "max_tokens" in params + assert "verbosity" not in params + + def test_verbosity_passes_through_for_gpt_5_models(self): + cfg = BedrockMantleChatConfig() + assert "verbosity" in cfg.get_supported_openai_params("openai.gpt-5.6-sol") + optional_params = litellm.get_optional_params( + model="openai.gpt-5.6-sol", + custom_llm_provider="bedrock_mantle", + verbosity="low", + drop_params=False, + ) + assert optional_params["verbosity"] == "low" class TestBedrockMantleChatAuth: From 884e96958c31c792a373b3e973b8cd20c338d3a9 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:27:48 +0000 Subject: [PATCH 176/187] fix(streaming): cap estimated reasoning tokens to the provider total and cover dict chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 20 ++--- .../test_streaming_chunk_builder_cursor.py | 81 +++++++++++++------ 2 files changed, 68 insertions(+), 33 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 99c02cdfc6a..9bb2d489f83 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -5,7 +5,6 @@ from itertools import groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast -from pydantic import BaseModel from typing_extensions import ReadOnly, Required from litellm._logging import verbose_logger @@ -974,12 +973,10 @@ class ChunkProcessor: return None @staticmethod - def _chunk_choices(chunk: "_UsageBearingChunk | BaseModel") -> Sequence[object]: + def _chunk_choices(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Sequence[object]: if isinstance(chunk, dict): return chunk.get("choices", ()) - if isinstance(chunk, (ModelResponse, ModelResponseStream)): - return chunk.choices - return () + return getattr(chunk, "choices", ()) @staticmethod def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool: @@ -1099,19 +1096,22 @@ class ChunkProcessor: if reasoning_tokens is not None: if returned_usage.completion_tokens_details is None: + capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens) returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=reasoning_tokens, - text_tokens=max(0, returned_usage.completion_tokens - reasoning_tokens), + reasoning_tokens=capped_reasoning_tokens, + text_tokens=returned_usage.completion_tokens - capped_reasoning_tokens, ) elif ( returned_usage.completion_tokens_details is not None and returned_usage.completion_tokens_details.reasoning_tokens is None ): - capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens) - returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens + existing_capped_reasoning_tokens: Final = min( + max(0, reasoning_tokens), returned_usage.completion_tokens + ) + returned_usage.completion_tokens_details.reasoning_tokens = existing_capped_reasoning_tokens if returned_usage.completion_tokens_details.text_tokens is None: returned_usage.completion_tokens_details.text_tokens = ( - returned_usage.completion_tokens - capped_reasoning_tokens + returned_usage.completion_tokens - existing_capped_reasoning_tokens ) if prompt_tokens_details is not None: returned_usage.prompt_tokens_details = prompt_tokens_details diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py index e925fd9b4a8..8617c5b81e8 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py @@ -19,7 +19,6 @@ to 0 when the only update we saw was the cursor, allowing the text-based fallback to estimate from the real completion text. """ - import pytest import litellm @@ -71,9 +70,7 @@ class TestAnthropicCursorBug: token_counter fallback can estimate from completion text. """ # Anthropic message_start: input_tokens accurate, output_tokens=1 cursor - message_start = _make_chunk( - usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) - ) + message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)) # Several content_block_delta chunks (no usage attached) text_chunks = [ _make_chunk(content="Hello"), @@ -99,9 +96,7 @@ class TestAnthropicCursorBug: Normal complete stream: message_start cursor=1, then message_delta=3847. Last-wins must give 3847 (the real value). """ - message_start = _make_chunk( - usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) - ) + message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)) text_chunks = [_make_chunk(content=t) for t in ["Hello", " world", "!"]] # message_delta with the real cumulative output_tokens message_delta = _make_chunk( @@ -121,19 +116,14 @@ class TestAnthropicCursorBug: End-to-end via calculate_usage(): cursor-only stream + real completion text should produce a token-counter estimate, NOT 1. """ - message_start = _make_chunk( - usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) - ) + message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)) # ~50 visible chars ≈ ~12 tokens (anthropic-style tokenizer ballpark) text_chunks = [ _make_chunk(content="Based on your question, I think the answer is "), _make_chunk(content="forty-two. Here is my reasoning: "), ] chunks = [message_start, *text_chunks] - completion_output = ( - "Based on your question, I think the answer is forty-two. " - "Here is my reasoning: " - ) + completion_output = "Based on your question, I think the answer is forty-two. Here is my reasoning: " processor = ChunkProcessor(chunks=chunks, messages=[]) usage = processor.calculate_usage( @@ -151,9 +141,7 @@ class TestAnthropicCursorBug: def test_cache_fields_preserved_from_message_start(self): """cache_read / cache_creation come from message_start and must survive.""" - message_start_usage = Usage( - prompt_tokens=1024, completion_tokens=1, total_tokens=1025 - ) + message_start_usage = Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) # Anthropic puts these in message_start message_start_usage.cache_read_input_tokens = 512 message_start_usage.cache_creation_input_tokens = 128 @@ -195,9 +183,7 @@ class TestAnthropicCursorBug: on a 1-token string also gives ~1, so billing is still approximately correct. This test pins that the result is sane (1 or 0). """ - message_start = _make_chunk( - usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21) - ) + message_start = _make_chunk(usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21)) text_chunk = _make_chunk(content="Yes.") # Anthropic's message_delta also gives output_tokens=1 in this case message_delta = _make_chunk( @@ -233,9 +219,7 @@ class TestAnthropicCursorBug: must fire so token_counter estimates from completion text instead of billing the placeholder. """ - message_start_usage = Usage( - prompt_tokens=1024, completion_tokens=1, total_tokens=1025 - ) + message_start_usage = Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025) message_start_usage.cache_read_input_tokens = 4096 message_start = _make_chunk(usage=message_start_usage) # Subsequent chunks with cache fields but no completion_tokens @@ -312,6 +296,57 @@ class TestAnthropicCursorBug: result = processor._calculate_usage_per_chunk(chunks=chunks) assert result["completion_tokens"] == 5 + def test_dict_chunks_with_finish_reason_are_trusted(self): + chunks = [ + { + "_hidden_params": {"custom_llm_provider": "anthropic"}, + "choices": [{"delta": {"content": "Yes, "}, "finish_reason": None}], + }, + { + "_hidden_params": {"custom_llm_provider": "anthropic"}, + "choices": [{"delta": {"content": "that works."}, "finish_reason": "stop"}], + "usage": Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25), + }, + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 5 + + def test_dict_chunks_without_finish_reason_reset_placeholder(self): + chunks = [ + { + "_hidden_params": {"custom_llm_provider": "anthropic"}, + "choices": [], + "usage": Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21), + }, + { + "_hidden_params": {"custom_llm_provider": "anthropic"}, + "choices": [{"delta": {"content": "partial"}, "finish_reason": None}], + }, + ] + processor = ChunkProcessor(chunks=chunks, messages=[]) + result = processor._calculate_usage_per_chunk(chunks=chunks) + assert result["completion_tokens"] == 0 + assert result["completion_tokens_details"] is None + + def test_estimated_reasoning_is_capped_to_trusted_completion_total(self): + chunks = [ + _make_chunk(reasoning_content="Let me reason about this carefully and at length. " * 20), + _make_chunk( + finish_reason="stop", + usage=Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25), + ), + ] + response = litellm.stream_chunk_builder( + chunks=chunks, + messages=[{"role": "user", "content": "Go."}], + ) + details = response.usage.completion_tokens_details + assert response.usage.completion_tokens == 5 + assert details.reasoning_tokens <= response.usage.completion_tokens + assert details.reasoning_tokens + details.text_tokens == response.usage.completion_tokens + assert details.text_tokens >= 0 + class TestProviderGuard: """Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic From aea9678f6198002fc6f914733320d921a5aff558 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:28:28 +0000 Subject: [PATCH 177/187] refactor(streaming): compute the reasoning token cap once Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/streaming_chunk_builder_utils.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 9bb2d489f83..0ca93fe08b3 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1095,8 +1095,8 @@ class ChunkProcessor: returned_usage.completion_tokens_details = completion_tokens_details if reasoning_tokens is not None: + capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens) if returned_usage.completion_tokens_details is None: - capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens) returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( reasoning_tokens=capped_reasoning_tokens, text_tokens=returned_usage.completion_tokens - capped_reasoning_tokens, @@ -1105,13 +1105,10 @@ class ChunkProcessor: returned_usage.completion_tokens_details is not None and returned_usage.completion_tokens_details.reasoning_tokens is None ): - existing_capped_reasoning_tokens: Final = min( - max(0, reasoning_tokens), returned_usage.completion_tokens - ) - returned_usage.completion_tokens_details.reasoning_tokens = existing_capped_reasoning_tokens + returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens if returned_usage.completion_tokens_details.text_tokens is None: returned_usage.completion_tokens_details.text_tokens = ( - returned_usage.completion_tokens - existing_capped_reasoning_tokens + returned_usage.completion_tokens - capped_reasoning_tokens ) if prompt_tokens_details is not None: returned_usage.prompt_tokens_details = prompt_tokens_details From e8f246bb6b36ba5e6d88bcbf3d43ba2a12082c84 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:32:32 +0000 Subject: [PATCH 178/187] fix(responses_bridge): forward verbosity as text.verbosity Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 6 +++- ...responses_transformation_transformation.py | 33 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 5a6debc4af5..a76032854b7 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -502,7 +502,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): elif key == "response_format": text_format = self._transform_response_format_to_text_format(value) if text_format: - responses_api_request["text"] = text_format + existing_text = cast("dict[str, object]", responses_api_request["text"] if "text" in responses_api_request else {}) + responses_api_request["text"] = cast("ResponseText", {**existing_text, **text_format}) + elif key == "verbosity": + existing_text = cast("dict[str, object]", responses_api_request["text"] if "text" in responses_api_request else {}) + responses_api_request["text"] = cast("ResponseText", {**existing_text, "verbosity": value}) elif key == "tool_choice": responses_api_request["tool_choice"] = self._normalize_tool_choice_for_responses_api(value) elif key == "stream_options": diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index e4ada0a9b31..c326ad4a0f7 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -4287,3 +4287,36 @@ def test_system_string_after_a_developer_message_stays_in_input_in_client_order( assert instructions is None assert [item["role"] for item in input_items] == ["developer", "system", "user"] assert input_items[1] == _system_input_item("Be brief.") + + +def test_map_optional_params_verbosity_merges_into_text(): + """Chat verbosity must land on Responses text.verbosity alongside text.format regardless of key order.""" + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams + + handler: Final = LiteLLMResponsesTransformationHandler() + + responses_api_request = ResponsesAPIOptionalRequestParams() + handler._map_optional_params_to_responses_api_request( + {"verbosity": "low", "response_format": {"type": "json_object"}}, + responses_api_request, + ) + assert responses_api_request["text"]["verbosity"] == "low" + assert responses_api_request["text"]["format"]["type"] == "json_object" + + reversed_request = ResponsesAPIOptionalRequestParams() + handler._map_optional_params_to_responses_api_request( + {"response_format": {"type": "json_object"}, "verbosity": "low"}, + reversed_request, + ) + assert reversed_request["text"]["verbosity"] == "low" + assert reversed_request["text"]["format"]["type"] == "json_object" + + verbosity_only_request = ResponsesAPIOptionalRequestParams() + handler._map_optional_params_to_responses_api_request( + {"verbosity": "low"}, + verbosity_only_request, + ) + assert verbosity_only_request["text"] == {"verbosity": "low"} From 438681be1af44cc5b846736a8a226a4083c4421f Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:33:28 +0000 Subject: [PATCH 179/187] refactor(responses_bridge): share text merge helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index a76032854b7..e1d915dca39 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -502,11 +502,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): elif key == "response_format": text_format = self._transform_response_format_to_text_format(value) if text_format: - existing_text = cast("dict[str, object]", responses_api_request["text"] if "text" in responses_api_request else {}) - responses_api_request["text"] = cast("ResponseText", {**existing_text, **text_format}) + responses_api_request["text"] = self._merge_text(responses_api_request, text_format) elif key == "verbosity": - existing_text = cast("dict[str, object]", responses_api_request["text"] if "text" in responses_api_request else {}) - responses_api_request["text"] = cast("ResponseText", {**existing_text, "verbosity": value}) + responses_api_request["text"] = self._merge_text( + responses_api_request, {"verbosity": cast(object, value)} + ) elif key == "tool_choice": responses_api_request["tool_choice"] = self._normalize_tool_choice_for_responses_api(value) elif key == "stream_options": @@ -522,6 +522,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): elif key == "web_search_options": self._add_web_search_tool(responses_api_request, value) + @staticmethod + def _merge_text( + responses_api_request: "ResponsesAPIOptionalRequestParams", update: Mapping[str, object] + ) -> "ResponseText": + existing: Final = cast( + "dict[str, object]", responses_api_request["text"] if "text" in responses_api_request else {} + ) + return cast("ResponseText", {**existing, **update}) + def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]: """Build sanitized litellm_params with merged metadata.""" responses_optional_param_keys: Final = set(ResponsesAPIOptionalRequestParams.__annotations__.keys()) From 41737aeda8d919b82c8018c2789e29211a37c4be Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:37:32 -0700 Subject: [PATCH 180/187] test(proxy): drop the docstring from the agent-runtime passthrough regression class --- .../pass_through_endpoints/test_llm_pass_through_endpoints.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 7f044486e14..6e82c90514d 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1984,8 +1984,6 @@ class TestBedrockAgentRuntimePassthroughToggle: class TestBedrockAgentRuntimePassthroughVirtualKeyLeak: - """Regression for LIT-7912: the agent-runtime branch of ``/bedrock/{endpoint}`` forwarded every caller header, - so a LiteLLM key presented in ``x-api-key`` or ``x-litellm-api-key`` rode to AWS next to the SigV4 signature.""" VKEY: Final = "sk-litellm-victim-key" MASTER_KEY: Final = "sk-master-1234" From 3ee8d43fdd63e068ee876a1999cf6412265fe49c Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 16:37:35 -0700 Subject: [PATCH 181/187] fix(ui): send only a changed TPM limit from the team admin settings form Save stays disabled until the value differs from the team's, so an unchanged form never reaches /team/update --- .../team/TeamAdminSettingsForm.test.tsx | 13 ++++++++++++ .../components/team/TeamAdminSettingsForm.tsx | 7 +++++-- .../team/teamAdminEditAccess.test.ts | 20 ++++++++++++++----- .../components/team/teamAdminEditAccess.ts | 8 +++++++- 4 files changed, 40 insertions(+), 8 deletions(-) diff --git a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx index 65f7a13652d..677c8859eb2 100644 --- a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx @@ -53,6 +53,18 @@ describe("TeamAdminSettingsForm", () => { await waitFor(() => expect(onSave).toHaveBeenCalledWith({ tpm_limit: null })); }); + it("keeps Save disabled until the TPM limit differs from the team's", () => { + renderForm(new Set(["tpm_limit"])); + const tpmInput = screen.getByLabelText("Tokens per minute Limit (TPM)"); + const save = screen.getByRole("button", { name: /save changes/i }); + + expect(save).toBeDisabled(); + fireEvent.change(tpmInput, { target: { value: "5000" } }); + expect(save).toBeEnabled(); + fireEvent.change(tpmInput, { target: { value: "1000" } }); + expect(save).toBeDisabled(); + }); + it("closes without saving on cancel", async () => { const user = userEvent.setup(); const { onSave, onCancel } = renderForm(new Set(["tpm_limit"])); @@ -65,6 +77,7 @@ describe("TeamAdminSettingsForm", () => { it("locks both buttons while a save is in flight", () => { renderForm(new Set(["tpm_limit"]), { isSaving: true }); + fireEvent.change(screen.getByLabelText("Tokens per minute Limit (TPM)"), { target: { value: "5000" } }); expect(screen.getByRole("button", { name: "Cancel" })).toBeDisabled(); expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled(); diff --git a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx index 581215b7fdf..140533fada5 100644 --- a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx @@ -1,6 +1,7 @@ "use client"; import { Save } from "lucide-react"; +import { useWatch } from "react-hook-form"; import { z } from "zod/v4"; import { FormField } from "@/components/shared/form/FormField"; @@ -37,7 +38,9 @@ export default function TeamAdminSettingsForm({ onSave, }: TeamAdminSettingsFormProps) { const form = useZodForm(teamAdminSettingsSchema, { defaultValues: initialValues }); - const submit = form.handleSubmit((values) => onSave(teamAdminSettingsChanges(values, editableFields))); + const draft = useWatch({ control: form.control }); + const hasChanges = Object.keys(teamAdminSettingsChanges(draft, initialValues, editableFields)).length > 0; + const submit = form.handleSubmit((values) => onSave(teamAdminSettingsChanges(values, initialValues, editableFields))); return (
void submit(event)}> @@ -56,7 +59,7 @@ export default function TeamAdminSettingsForm({ - diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts index c8117800053..ded6d775ce8 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts @@ -20,21 +20,31 @@ describe("teamAdminFieldLabel", () => { describe("teamAdminSettingsChanges", () => { const tpmEnabled = new Set(["tpm_limit"]); + const stored = { tpm_limit: 1000 }; it.each([ ["a typed number string", "5000", 5000], - ["a stored number", 1200, 1200], + ["a number", 1200, 1200], ["zero", "0", 0], ["an emptied input", "", null], ["whitespace", " ", null], - ["no stored limit", null, null], + ["no limit", null, null], ["an unset value", undefined, null], - ])("sends tpm_limit for %s", (_label, tpm_limit, expected) => { - expect(teamAdminSettingsChanges({ tpm_limit }, tpmEnabled)).toStrictEqual({ tpm_limit: expected }); + ])("sends tpm_limit changed to %s", (_label, tpm_limit, expected) => { + expect(teamAdminSettingsChanges({ tpm_limit }, stored, tpmEnabled)).toStrictEqual({ tpm_limit: expected }); + }); + + it.each([ + ["the stored number", 1000, { tpm_limit: 1000 }], + ["the stored number typed back in", "1000", { tpm_limit: 1000 }], + ["an emptied input over no stored limit", "", { tpm_limit: null }], + ["an unset value over no stored limit", undefined, { tpm_limit: null }], + ])("sends nothing for %s", (_label, tpm_limit, initialValues) => { + expect(teamAdminSettingsChanges({ tpm_limit }, initialValues, tpmEnabled)).toStrictEqual({}); }); it("leaves tpm_limit out when the proxy did not enable it for team admins", () => { - expect(teamAdminSettingsChanges({ tpm_limit: "5000" }, new Set(["max_budget"]))).toStrictEqual({}); + expect(teamAdminSettingsChanges({ tpm_limit: "5000" }, stored, new Set(["max_budget"]))).toStrictEqual({}); }); }); diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts index d38566eefda..73129923907 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts @@ -59,8 +59,14 @@ const numberOrNull = (value: string | number | null | undefined): number | null export const teamAdminSettingsChanges = ( values: TeamAdminSettingsValues, + initialValues: TeamAdminSettingsValues, editableFields: ReadonlySet, -): TeamAdminSettingsChanges => (editableFields.has("tpm_limit") ? { tpm_limit: numberOrNull(values.tpm_limit) } : {}); +): TeamAdminSettingsChanges => { + const tpmLimit = numberOrNull(values.tpm_limit); + return editableFields.has("tpm_limit") && tpmLimit !== numberOrNull(initialValues.tpm_limit) + ? { tpm_limit: tpmLimit } + : {}; +}; export const parseTeamEditAccess = (callerEditAccess: unknown): TeamEditAccess => { const parsed = callerEditAccessSchema.safeParse(callerEditAccess); From 255ef3bc267bea22259415c7f79a1f135a46de35 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:37:45 +0000 Subject: [PATCH 182/187] refactor(bedrock_mantle): build supported params without mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../bedrock_mantle/chat/transformation.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 5b69d7aff42..41d93a8dd4d 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -109,15 +109,22 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): def get_supported_openai_params(self, model: str) -> list: base_params: Final = super().get_supported_openai_params(model) - if is_gpt_reasoning_series_name(model) and "verbosity" not in base_params: - base_params.append("verbosity") + extra_params: Final = tuple( + param + for param, supported in ( + ("verbosity", is_gpt_reasoning_series_name(model)), + ("reasoning_effort", self._supports_reasoning(model)), + ) + if supported and param not in base_params + ) + return [*base_params, *extra_params] + + def _supports_reasoning(self, model: str) -> bool: try: - if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): - if "reasoning_effort" not in base_params: - base_params.append("reasoning_effort") + return litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider) except Exception as e: verbose_logger.debug("BedrockMantleChatConfig: error checking reasoning support: %s", e) - return base_params + return False def get_model_response_iterator( self, From 92121242666fade26b8680063bdb9e5a749e4ce5 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:37:45 +0000 Subject: [PATCH 183/187] style(responses_bridge): use dict.get in text merge helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_responses_transformation/transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index e1d915dca39..eaf7f47552c 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -526,9 +526,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _merge_text( responses_api_request: "ResponsesAPIOptionalRequestParams", update: Mapping[str, object] ) -> "ResponseText": - existing: Final = cast( - "dict[str, object]", responses_api_request["text"] if "text" in responses_api_request else {} - ) + existing: Final = cast("dict[str, object]", dict(responses_api_request).get("text") or {}) return cast("ResponseText", {**existing, **update}) def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]: From 574ea15b8fa72480419e3d4d76c74ad38d920e62 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:42:20 -0700 Subject: [PATCH 184/187] fix(mcp): count admin static headers as api_key credential slots An api_key server whose key lives in static_headers, the documented shape for upstreams that expect a custom header name, dispatched fine before the fail-closed check and was rejected as misconfigured after it. The check now treats every static header the admin configured as a credential slot for api_key mode, on both the MCP client path and the OpenAPI tool path, with regression tests at all three layers. --- .../mcp_server/openapi_to_mcp_generator.py | 2 +- .../outbound_credentials/adapter.py | 9 ++++++-- .../outbound_credentials/test_adapter.py | 18 ++++++++++++++- .../mcp_server/test_mcp_server_manager.py | 22 +++++++++++++++++++ .../test_openapi_to_mcp_generator.py | 20 +++++++++++++++++ 5 files changed, 67 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 477d86ab436..0cdf40ae8d3 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -512,7 +512,7 @@ def create_tool_function( ) from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok - match validate_static_credential(auth_type, effective_headers, upstream_token_header): + match validate_static_credential(auth_type, effective_headers, upstream_token_header, headers or ()): case Error(error): raise_public(error) case Ok(): diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index ba223f73b2d..42947e39530 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -13,7 +13,7 @@ from __future__ import annotations import base64 import os -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Final, Literal, NoReturn from fastapi import HTTPException @@ -426,16 +426,19 @@ def validate_static_credential( auth_type: MCPAuthType, headers: Mapping[str, str], upstream_token_header: str | None = None, + static_header_names: Iterable[str] = (), ) -> Result[None, CredError]: if auth_type not in _STATIC_MODES: return Ok(None) default_slot: Final = "X-API-Key" if auth_type == MCPAuth.api_key else "Authorization" + admin_chosen_slots: Final = tuple(static_header_names) if auth_type == MCPAuth.api_key else () slots: Final = frozenset( name.lower() for name in ( upstream_token_header or default_slot, default_slot, "Authorization", + *admin_chosen_slots, ) ) values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) @@ -448,7 +451,9 @@ async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient: if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio: return client request: Final = await client.prepare_request_auth() - match validate_static_credential(server.auth_type, request.headers, server.upstream_token_header): + match validate_static_credential( + server.auth_type, request.headers, server.upstream_token_header, server.static_headers or () + ): case Error(error): raise_public(error) case Ok(): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index 78da9ff4d77..141260db700 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -23,7 +23,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import to_subject, validate_static_credential, ) -from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( ApiKeyConfig, AuthorizationCodeConfig, @@ -59,6 +59,22 @@ def test_static_credential_preserves_supported_api_key_and_raw_headers( assert isinstance(result, Ok) +@pytest.mark.parametrize("auth_type,headers,static_header_names,expected", [ + (MCPAuth.api_key, {"apikey": "static-key"}, ("apikey",), Ok), + (MCPAuth.api_key, {"apikey": "static-key", "X-API-Key": ""}, ("apikey",), Ok), + (MCPAuth.api_key, {"apikey": ""}, ("apikey",), Error), + (MCPAuth.api_key, {"apikey": "static-key"}, (), Error), + (MCPAuth.api_key, {"apikey": "static-key"}, ("X-Tenant",), Error), + (MCPAuth.bearer_token, {"apikey": "static-key"}, ("apikey",), Error), + (MCPAuth.token, {"apikey": "static-key"}, ("apikey",), Error), +]) +def test_static_credential_counts_api_key_static_headers_only( + auth_type: MCPAuthType, headers: dict[str, str], static_header_names: tuple[str, ...], expected: type, +) -> None: + result: Final = validate_static_credential(auth_type, headers, static_header_names=static_header_names) + assert isinstance(result, expected) + + def _server(**kwargs) -> MCPServer: return MCPServer(server_id="s", name="n", transport=MCPTransport.http, **kwargs) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 186c46e1b37..2fab7a6f4b5 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 @@ -13652,6 +13652,28 @@ class TestProtectedCredentialPreparation: assert client._credential_slot == "X-Custom" assert await client.discovery_auth_fingerprint() + @pytest.mark.asyncio + @pytest.mark.parametrize("static_headers,accepted", [ + ({"apikey": "static-key"}, True), + ({"apikey": ""}, False), + ({"X-Tenant": "tenant"}, True), + ]) + async def test_api_key_carried_by_static_header_passes_fail_closed_check( + self, static_headers: dict[str, str], accepted: bool + ) -> None: + server: Final = MCPServer( + server_id="static-slot", name="static-slot", url="https://upstream.example/mcp", + transport=MCPTransport.http, auth_type=MCPAuth.api_key, static_headers=static_headers, + ) + if not accepted: + with pytest.raises(HTTPException) as exc: + await MCPServerManager()._create_mcp_client(server, extra_headers=dict(static_headers)) + assert exc.value.status_code == 500 + return + client: Final = await MCPServerManager()._create_mcp_client(server, extra_headers=dict(static_headers)) + request: Final = await client.prepare_request_auth() + assert all(request.headers[name] == value for name, value in static_headers.items()) + @pytest.mark.asyncio @pytest.mark.parametrize("static,forwarded,caller", [ ({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index a9def20e75d..bd351f9106e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -133,6 +133,26 @@ async def test_static_auth_uses_configured_custom_header( assert destination.call_count == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("credential", ["static-key", ""]) +async def test_static_auth_accepts_api_key_carried_by_static_header( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, credential: str, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", headers={"apikey": credential}, auth_type=MCPAuth.api_key, + ) + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + if credential: + assert await tool() == "authenticated" + assert destination.calls.last.request.headers["apikey"] == credential + assert "x-api-key" not in destination.calls.last.request.headers + else: + with pytest.raises(HTTPException, match="requires a usable upstream credential"): + await tool() + assert destination.call_count == 0 + + @pytest.mark.asyncio @pytest.mark.parametrize("auth_type,resolved", [ (MCPAuth.none, None), From b7a9042f1b2930303dd67ec316c43e48348b7909 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:43:04 +0000 Subject: [PATCH 185/187] style(responses_bridge): suppress type-discipline flags with reasons Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../transformation.py | 12 +++++++++--- litellm/llms/bedrock_mantle/chat/transformation.py | 2 +- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index eaf7f47552c..11fa2a2bab6 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -505,7 +505,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): responses_api_request["text"] = self._merge_text(responses_api_request, text_format) elif key == "verbosity": responses_api_request["text"] = self._merge_text( - responses_api_request, {"verbosity": cast(object, value)} + responses_api_request, + MappingProxyType({"verbosity": value}), # pyright: ignore[reportUnknownArgumentType] # untyped value ) elif key == "tool_choice": responses_api_request["tool_choice"] = self._normalize_tool_choice_for_responses_api(value) @@ -526,8 +527,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _merge_text( responses_api_request: "ResponsesAPIOptionalRequestParams", update: Mapping[str, object] ) -> "ResponseText": - existing: Final = cast("dict[str, object]", dict(responses_api_request).get("text") or {}) - return cast("ResponseText", {**existing, **update}) + existing: Final = cast( # cast-ok: text field is a ResponseText | dict[str, Any] | None union + "dict[str, object]", + dict(responses_api_request).get("text") or {}, # mutable-ok: one-shot merge seed + ) + return cast( # cast-ok: merged mapping is a valid ResponseText shape + "ResponseText", {**existing, **update} # mutable-ok: one-shot merged payload + ) def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]: """Build sanitized litellm_params with merged metadata.""" diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 41d93a8dd4d..590919f1fb0 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -117,7 +117,7 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): ) if supported and param not in base_params ) - return [*base_params, *extra_params] + return [*base_params, *extra_params] # mutable-ok: fresh list required by the inherited signature def _supports_reasoning(self, model: str) -> bool: try: From ada0a1ad3a06d3ae970a0e5e61d224867fb2fcf5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:49:29 -0700 Subject: [PATCH 186/187] fix(azure_ai): strip the azure_ai/ prefix when a Responses call is remapped to azure A catalog OpenAI name on an .openai.azure.com host (or with AZURE_AI_API_BASE set to one) is remapped from azure_ai to azure before the Responses request is built, and the azure_ai/ prefix stayed in the wire model, so Azure answered DeploymentNotFound. The Azure Responses config now strips azure_ai/ next to responses/ and o_series/. --- litellm/llms/azure/responses/transformation.py | 7 +------ .../response/test_azure_transformation.py | 11 +++++++++++ .../test_azure_ai_responses_transformation.py | 18 ++++++++++++++++++ 3 files changed, 30 insertions(+), 6 deletions(-) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 7fe12138ebc..2a82b42df7b 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -49,12 +49,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params) def get_stripped_model_name(self, model: str) -> str: - # if "responses/" is in the model name, remove it - if "responses/" in model: - model = model.replace("responses/", "") - if "o_series" in model: - model = model.replace("o_series/", "") - return model + return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "") def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]: """ diff --git a/tests/test_litellm/llms/azure/response/test_azure_transformation.py b/tests/test_litellm/llms/azure/response/test_azure_transformation.py index 726c9f65681..532c278e891 100644 --- a/tests/test_litellm/llms/azure/response/test_azure_transformation.py +++ b/tests/test_litellm/llms/azure/response/test_azure_transformation.py @@ -677,3 +677,14 @@ def test_azure_responses_gpt6_astra_rejects_temperature_while_reasoning(local_mo model="gpt-6-astra", drop_params=False, ) + + +def test_azure_responses_sends_the_deployment_name_when_azure_ai_prefix_survives_provider_remap(): + request = AzureOpenAIResponsesAPIConfig().transform_responses_api_request( + model="azure_ai/gpt-5.4-nano", + input="hi", + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert request["model"] == "gpt-5.4-nano" diff --git a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py index 925608a3c9b..bae956eb061 100644 --- a/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py +++ b/tests/test_litellm/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -253,6 +253,24 @@ async def test_aresponses_sends_reasoning_and_tools_to_native_endpoint(model, ap _assert_native_responses_request(route, expected_url, expected_model) +@pytest.mark.asyncio +@respx.mock +async def test_aresponses_catalog_name_remapped_to_azure_sends_bare_deployment_name(monkeypatch): + monkeypatch.setenv("AZURE_AI_API_BASE", "https://res.openai.azure.com") + route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock( + return_value=httpx.Response(200, json=_responses_payload("gpt-5.4-nano")) + ) + + await litellm.aresponses( + model="azure_ai/gpt-5.4-nano", + input="What is the weather in SF?", + api_base="https://res.openai.azure.com", + api_key="fake-key", + ) + + assert json.loads(route.calls.last.request.content)["model"] == "gpt-5.4-nano" + + @pytest.mark.asyncio @respx.mock @pytest.mark.parametrize("model,api_base,expected_url,expected_model", NATIVE_RESPONSES_CASES) From bba15b13824d392bee91bf0e36ede1d99e84fc40 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:52:53 +0000 Subject: [PATCH 187/187] style(responses_bridge): apply ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_responses_transformation/transformation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 11fa2a2bab6..1b976f5a48b 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -532,7 +532,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): dict(responses_api_request).get("text") or {}, # mutable-ok: one-shot merge seed ) return cast( # cast-ok: merged mapping is a valid ResponseText shape - "ResponseText", {**existing, **update} # mutable-ok: one-shot merged payload + "ResponseText", + {**existing, **update}, # mutable-ok: one-shot merged payload ) def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]: