From 26805f75fe44170f36b54d4e8e57d9ce7ad19b5b Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 5 Aug 2026 19:36:37 +0000 Subject: [PATCH 01/26] fix(bedrock_mantle): stop dropping the web_search tool on /v1/responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../responses/transformation.py | 4 +- ...odel_prices_and_context_window_backup.json | 15 ++-- model_prices_and_context_window.json | 15 ++-- ...bedrock_mantle_responses_transformation.py | 90 +++++++++++++++++-- 4 files changed, 107 insertions(+), 17 deletions(-) diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 92da5835b2d..3b01c7dbad0 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -44,7 +44,9 @@ _BASE_SUFFIXES_TO_STRIP: Final = ( ) # Per Bedrock Mantle Responses API validation errors. -_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES = frozenset({"function", "mcp", "custom", "namespace", "tool_search"}) +_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset( + {"function", "mcp", "custom", "namespace", "tool_search", "web_search"} +) _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"}) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 02b3cde217a..e85d49cfd12 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45247,7 +45247,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-terra": { "input_cost_per_token": 2.2e-06, @@ -45275,7 +45276,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, @@ -45303,7 +45305,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, @@ -45330,7 +45333,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.4": { "input_cost_per_token": 2.75e-06, @@ -45357,7 +45361,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/google.gemma-4-31b": { "input_cost_per_token": 1.4e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bc7330ec99c..7514f1596b0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45368,7 +45368,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-terra": { "input_cost_per_token": 2.2e-06, @@ -45396,7 +45397,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, @@ -45424,7 +45426,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, @@ -45451,7 +45454,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.4": { "input_cost_per_token": 2.75e-06, @@ -45478,7 +45482,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/google.gemma-4-31b": { "input_cost_per_token": 1.4e-07, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index bea979aec64..3810ef39062 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -339,7 +339,7 @@ class TestBedrockMantleResponsesTools: params = cfg.map_openai_params( response_api_optional_params={ "tools": [ - {"type": "web_search"}, + {"type": "file_search", "vector_store_ids": ["vs_123"]}, {"type": "function", "name": "exec_command"}, ] }, @@ -351,7 +351,7 @@ class TestBedrockMantleResponsesTools: def test_map_openai_params_removes_tools_when_all_unsupported(self): cfg = BedrockMantleResponsesAPIConfig() params = cfg.map_openai_params( - response_api_optional_params={"tools": [{"type": "web_search"}]}, + response_api_optional_params={"tools": [{"type": "file_search", "vector_store_ids": ["vs_123"]}]}, model="openai.gpt-5.5", drop_params=False, ) @@ -365,12 +365,86 @@ class TestBedrockMantleResponsesTools: "litellm.llms.bedrock_mantle.responses.transformation.verbose_logger.warning" ) as mock_warning: cfg.map_openai_params( - response_api_optional_params={"tools": [{"type": "web_search"}]}, + response_api_optional_params={"tools": [{"type": "file_search", "vector_store_ids": ["vs_123"]}]}, model="openai.gpt-5.5", drop_params=False, ) assert mock_warning.call_count == 1 - assert "web_search" in str(mock_warning.call_args) + assert "file_search" in str(mock_warning.call_args) + + +class TestBedrockMantleResponsesWebSearch: + """Web Search on Amazon Bedrock is a server-side built-in tool that Mantle runs + itself when the caller passes {"type": "web_search"} on the Responses path, so + the config must forward the tool and its options untouched instead of filtering + it out and returning an ungrounded answer.""" + + _WEB_SEARCH_TOOL = {"type": "web_search", "external_web_access": False} + + def test_web_search_survives_map_openai_params_with_its_options(self): + cfg = BedrockMantleResponsesAPIConfig() + params = cfg.map_openai_params( + response_api_optional_params={"tools": [self._WEB_SEARCH_TOOL]}, + model="openai.gpt-5.6-sol", + drop_params=False, + ) + assert params["tools"] == [self._WEB_SEARCH_TOOL] + + def test_web_search_reaches_outbound_body_alongside_function_tools(self): + cfg = BedrockMantleResponsesAPIConfig() + function_tool = {"type": "function", "name": "exec_command"} + params = cfg.map_openai_params( + response_api_optional_params={"tools": [self._WEB_SEARCH_TOOL, function_tool]}, + model="openai.gpt-5.6-sol", + drop_params=False, + ) + body = cfg.transform_responses_api_request( + model="openai.gpt-5.6-sol", + input="What did AWS announce today?", + response_api_optional_request_params=params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["tools"] == [self._WEB_SEARCH_TOOL, function_tool] + + def test_web_search_is_not_logged_as_dropped(self): + from unittest.mock import patch + + cfg = BedrockMantleResponsesAPIConfig() + with patch("litellm.llms.bedrock_mantle.responses.transformation.verbose_logger.warning") as mock_warning: + cfg.map_openai_params( + response_api_optional_params={"tools": [self._WEB_SEARCH_TOOL]}, + model="openai.gpt-5.6-sol", + drop_params=False, + ) + assert mock_warning.call_count == 0 + + def test_hoisted_web_search_tool_survives(self): + cfg = BedrockMantleResponsesAPIConfig() + body = cfg.transform_responses_api_request( + model="openai.gpt-5.6-sol", + input=[ + {"type": "additional_tools", "role": "developer", "tools": [self._WEB_SEARCH_TOOL]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + ], + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["tools"] == [self._WEB_SEARCH_TOOL] + + @pytest.mark.parametrize( + "model", + [ + "bedrock_mantle/openai.gpt-5.6-sol", + "bedrock_mantle/openai.gpt-5.6-terra", + "bedrock_mantle/openai.gpt-5.6-luna", + "bedrock_mantle/openai.gpt-5.5", + "bedrock_mantle/openai.gpt-5.4", + ], + ) + def test_cost_map_advertises_web_search_support(self, model): + assert litellm.supports_web_search(model=model) is True def _codex_exec_tool(): @@ -558,7 +632,7 @@ class TestBedrockMantleCodexAdditionalTools: "type": "additional_tools", "role": "developer", "tools": [ - {"type": "web_search"}, + {"type": "file_search", "vector_store_ids": ["vs_123"]}, {"type": "function", "name": "wait"}, ], }, @@ -570,7 +644,11 @@ class TestBedrockMantleCodexAdditionalTools: def test_item_stripped_even_when_no_hoisted_tool_survives(self): body = self._transform( input=[ - {"type": "additional_tools", "role": "developer", "tools": [{"type": "web_search"}]}, + { + "type": "additional_tools", + "role": "developer", + "tools": [{"type": "file_search", "vector_store_ids": ["vs_123"]}], + }, self._USER_MESSAGE, ] ) From e36e3d4a6f5b77329d53ac461ed5d3a081f0f36d Mon Sep 17 00:00:00 2001 From: eeshsaxena <139802361+eeshsaxena@users.noreply.github.com> Date: Sat, 8 Aug 2026 15:56:45 +0530 Subject: [PATCH 02/26] fix(proxy): raise (not return) proxy exception in get/delete credential endpoints --- litellm/proxy/credential_endpoints/endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 3b3e9692eda..106b43a722c 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -130,7 +130,7 @@ async def get_credentials( ] return {"success": True, "credentials": masked_credentials} except Exception as e: - return handle_exception_on_proxy(e) + raise handle_exception_on_proxy(e) @router.get( @@ -240,7 +240,7 @@ async def delete_credential( litellm.credential_list = [cred for cred in litellm.credential_list if cred.credential_name != credential_name] return {"success": True, "message": "Credential deleted successfully"} except Exception as e: - return handle_exception_on_proxy(e) + raise handle_exception_on_proxy(e) def update_db_credential( From e543ae39808f61bfa416ef0450c535f75fe6b798 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:27:26 +0000 Subject: [PATCH 03/26] fix(databricks): strip thinking_blocks and reasoning_content from outbound messages Databricks Model Serving validates assistant messages with additionalProperties=false, so replaying a thinking turn translated by the Anthropic Messages adapter 400s with 'messages.N.thinking_blocks: Extra inputs are not permitted'. Drop litellm's internal fields in DatabricksConfig._transform_messages via a shared common_utils helper. Resolves LIT-6762 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../prompt_templates/common_utils.py | 14 ++++++++ .../llms/databricks/chat/transformation.py | 3 ++ .../test_databricks_chat_transformation.py | 34 +++++++++++++++++++ 3 files changed, 51 insertions(+) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ff46440ff5c..d5b1cbb7274 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1554,6 +1554,20 @@ def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT: return cast(_MarkedT, marked) # cast-ok: same block shape as the input plus the marker key +LITELLM_INTERNAL_MESSAGE_FIELDS: Final = frozenset({"thinking_blocks", "reasoning_content", "provider_specific_fields"}) + + +def strip_litellm_internal_message_fields(message: AllMessageValues) -> AllMessageValues: + """Drop the fields litellm attaches to assistant messages (e.g. when translating Anthropic thinking + blocks) that OpenAI-compatible endpoints with strict schemas reject as extra inputs.""" + if LITELLM_INTERNAL_MESSAGE_FIELDS.isdisjoint(message): + return message + return cast( # cast-ok: same TypedDict minus internal keys + AllMessageValues, + {key: value for key, value in message.items() if key not in LITELLM_INTERNAL_MESSAGE_FIELDS}, + ) + + def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any: """ Filters a value from a dictionary diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index c587146005f..e59db2dac19 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( + strip_litellm_internal_message_fields, strip_name_from_message, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -419,6 +420,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): """ Databricks does not support: - 'name' in user message. + - litellm's internal `thinking_blocks` / `reasoning_content` on assistant messages. """ new_messages = [] for idx, message in enumerate(messages): @@ -427,6 +429,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): else: _message = message _message = strip_name_from_message(_message, allowed_name_roles=["user"]) + _message = strip_litellm_internal_message_fields(_message) # Move message-level cache_control into a content block when content is a string. if "cache_control" in _message and isinstance(_message.get("content"), str): _message = self._move_cache_control_into_string_content_block(_message) diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 41fb2589655..0a792a546e4 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -255,6 +255,40 @@ def test_transform_messages_sanitizes_empty_content(): assert result[1]["content"] == "Hi" +def test_transform_request_strips_thinking_blocks_and_reasoning_content(): + """Regression for LIT-6762: replaying an assistant turn that litellm decorated with + `thinking_blocks` / `reasoning_content` made Databricks 400 with + 'messages.N.thinking_blocks: Extra inputs are not permitted'.""" + config = DatabricksConfig() + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "Hello! How can I help?", + "thinking_blocks": [ + {"type": "thinking", "thinking": "greet briefly", "signature": "sig_abc", "cache_control": {}} + ], + "reasoning_content": "greet briefly", + "provider_specific_fields": {"foo": "bar"}, + }, + {"role": "user", "content": "thanks"}, + ] + + result = config.transform_request( + model="databricks-claude-opus-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + )["messages"] + + assert result[1] == {"role": "assistant", "content": "Hello! How can I help?"} + assert not any( + key in message for message in result for key in ("thinking_blocks", "reasoning_content", "provider_specific_fields") + ) + assert "thinking_blocks" in messages[1] + + def _parallel_tool_calls(): return [ { From 518506834a3adfe068d07cb567fd27448c9f1bc6 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:43:22 +0000 Subject: [PATCH 04/26] fix(databricks): drop assistant turns left empty after stripping thinking_blocks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/databricks/chat/transformation.py | 10 +++++ .../test_databricks_chat_transformation.py | 44 ++++++++++++++++++- 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index e59db2dac19..4f61a39f692 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -56,6 +56,14 @@ from ...openai_like.chat.transformation import OpenAILikeChatConfig from ..common_utils import DatabricksBase, DatabricksException +def _is_bare_assistant_message(message_dict: dict[str, Any]) -> bool: + """Databricks rejects assistant messages with neither content nor tool calls, e.g. a replayed + thinking-only turn once its `thinking_blocks` are stripped.""" + return message_dict.get("role") == "assistant" and not any( + message_dict.get(key) for key in ("content", "tool_calls", "function_call") + ) + + def _sanitize_empty_content(message_dict: dict[str, Any]) -> None: """ Remove or filter content so empty text blocks are not sent. @@ -434,6 +442,8 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): if "cache_control" in _message and isinstance(_message.get("content"), str): _message = self._move_cache_control_into_string_content_block(_message) _sanitize_empty_content(cast(dict[str, Any], _message)) + if _is_bare_assistant_message(cast(dict[str, Any], _message)): + continue new_messages.append(_message) if "claude" not in model: diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index 0a792a546e4..b2c52617e9d 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -284,11 +284,53 @@ def test_transform_request_strips_thinking_blocks_and_reasoning_content(): assert result[1] == {"role": "assistant", "content": "Hello! How can I help?"} assert not any( - key in message for message in result for key in ("thinking_blocks", "reasoning_content", "provider_specific_fields") + key in message + for message in result + for key in ("thinking_blocks", "reasoning_content", "provider_specific_fields") ) assert "thinking_blocks" in messages[1] +def test_transform_request_drops_thinking_only_assistant_turn_but_keeps_tool_call_turn(): + """A replayed thinking-only assistant turn has nothing left once `thinking_blocks` are stripped, so it must be + dropped instead of being sent as a bare {"role": "assistant"}. A thinking + tool_use turn keeps its tool_calls.""" + config = DatabricksConfig() + tool_call = {"id": "call_1", "type": "function", "function": {"name": "f", "arguments": "{}"}} + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "thinking_blocks": [{"type": "thinking", "thinking": "hmm", "signature": "sig_1"}], + "reasoning_content": "hmm", + }, + {"role": "user", "content": "again"}, + { + "role": "assistant", + "content": None, + "thinking_blocks": [{"type": "thinking", "thinking": "call f", "signature": "sig_2"}], + "reasoning_content": "call f", + "tool_calls": [tool_call], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "ok"}, + ] + + result = config.transform_request( + model="databricks-claude-opus-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + )["messages"] + + assert result == [ + {"role": "user", "content": "hi"}, + {"role": "user", "content": "again"}, + {"role": "assistant", "tool_calls": [tool_call]}, + {"role": "tool", "tool_call_id": "call_1", "content": "ok"}, + ] + + def _parallel_tool_calls(): return [ { From 1d40202a4e8863161f9a2339193d367895d7ea1a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 2 Sep 2026 21:48:25 +0000 Subject: [PATCH 05/26] fix(databricks): keep new helpers within the type discipline budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/prompt_templates/common_utils.py | 4 +++- litellm/llms/databricks/chat/transformation.py | 6 +++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index d5b1cbb7274..17ebde83eee 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1564,7 +1564,9 @@ def strip_litellm_internal_message_fields(message: AllMessageValues) -> AllMessa return message return cast( # cast-ok: same TypedDict minus internal keys AllMessageValues, - {key: value for key, value in message.items() if key not in LITELLM_INTERNAL_MESSAGE_FIELDS}, + { # mutable-ok: provider transforms mutate message dicts in place downstream + key: value for key, value in message.items() if key not in LITELLM_INTERNAL_MESSAGE_FIELDS + }, ) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 4f61a39f692..420c357ab87 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -3,7 +3,7 @@ Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completion """ import os -from collections.abc import AsyncIterator, Coroutine, Iterator +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload import httpx @@ -56,7 +56,7 @@ from ...openai_like.chat.transformation import OpenAILikeChatConfig from ..common_utils import DatabricksBase, DatabricksException -def _is_bare_assistant_message(message_dict: dict[str, Any]) -> bool: +def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool: """Databricks rejects assistant messages with neither content nor tool calls, e.g. a replayed thinking-only turn once its `thinking_blocks` are stripped.""" return message_dict.get("role") == "assistant" and not any( @@ -442,7 +442,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): if "cache_control" in _message and isinstance(_message.get("content"), str): _message = self._move_cache_control_into_string_content_block(_message) _sanitize_empty_content(cast(dict[str, Any], _message)) - if _is_bare_assistant_message(cast(dict[str, Any], _message)): + if _is_bare_assistant_message(_message): continue new_messages.append(_message) From 7b8cc0319e70d90dd096fbe86830b987a67ef491 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:56:09 -0700 Subject: [PATCH 06/26] fix(proxy-extras): only spend a migrate-deploy attempt when a pass made no progress The v2 migration resolver gave `prisma migrate deploy` four attempts, and every recovery path ended in a bare `continue`, so each one burned an attempt. A database first brought up with `--use_prisma_db_push` has a full schema and no migrations ledger, so the baseline spent attempt one and the first three migrations whose objects already existed spent the rest. The proxy then exited before binding its port, and that database could never be moved onto the resolver. The retry budget now counts only attempts that got nowhere. Creating the baseline, and each migration newly marked applied, leaves the budget alone, so a push-created database works through its pre-existing objects one pass at a time. Timeouts, deadlock rollbacks, advisory-lock waits, and a repeat of a recovery that already ran still spend an attempt, so a run that stops making progress gives up exactly as before. --- .../litellm_proxy_extras/utils.py | 67 ++++++- .../test_litellm_proxy_extras_utils.py | 167 ++++++++++++++++++ 2 files changed, 226 insertions(+), 8 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index d22484bc0e8..21f255b0666 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -6,6 +6,7 @@ import shutil import subprocess import tempfile import time +from dataclasses import dataclass, replace from pathlib import Path from typing import Optional @@ -45,6 +46,38 @@ _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") _MIGRATION_DEADLOCK_MARKER = "deadlock detected" +MAX_MIGRATE_DEPLOY_ATTEMPTS = 4 + + +@dataclass(frozen=True) +class _MigrateAttemptBudget: + """Retries left, and the recoveries already run. + + A recovery that lands something new costs nothing, so a database full of + objects `prisma db push` created works through them one per pass. Anything + that made no progress spends an attempt, so a stuck run still gives up. + """ + + attempts_left: int + recoveries: frozenset[str] = frozenset() + + @property + def exhausted(self) -> bool: + return self.attempts_left <= 0 + + @property + def attempt_number(self) -> int: + return MAX_MIGRATE_DEPLOY_ATTEMPTS - self.attempts_left + 1 + + def spend(self) -> "_MigrateAttemptBudget": + return replace(self, attempts_left=self.attempts_left - 1) + + def after_recovery(self, recovery: str) -> "_MigrateAttemptBudget": + if recovery in self.recoveries: + return self.spend() + return replace(self, recoveries=self.recoveries | {recovery}) + + _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE) _SPEND_LOGS_ARTIFACT_DROP_RE = re.compile( r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE @@ -716,6 +749,9 @@ class ProxyExtrasDBManager: Ahead-of-HEAD state (DB has migrations newer than this build ships) is logged as a warning, not a fatal error — users whose DBs got into weird shapes from the old thrashing should still be able to start. + + The retry budget only counts attempts that made no progress: see + _MigrateAttemptBudget. """ schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" migrations_dir = ProxyExtrasDBManager._get_prisma_dir() @@ -749,8 +785,9 @@ class ProxyExtrasDBManager: original_dir = os.getcwd() os.chdir(migrations_dir) deploy_timeout = prisma_migrate_deploy_timeout() + budget = _MigrateAttemptBudget(attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS) try: - for attempt in range(4): + while not budget.exhausted: try: result = subprocess.run( [_get_prisma_command(), "migrate", "deploy"], @@ -767,10 +804,11 @@ class ProxyExtrasDBManager: logger.warning( "prisma migrate deploy attempt %s timed out after %ss, retrying. " "Raise %s if this database needs longer to apply its pending migrations.", - attempt + 1, + budget.attempt_number, deploy_timeout, PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ) + budget = budget.spend() time.sleep(random.randrange(5, 15)) continue @@ -781,7 +819,14 @@ class ProxyExtrasDBManager: logger.info( "Schema exists but no migrations ledger — creating baseline" ) - ProxyExtrasDBManager._create_baseline_migration(schema_path) + baselined = ProxyExtrasDBManager._create_baseline_migration( + schema_path + ) + budget = ( + budget.after_recovery("baseline") + if baselined + else budget.spend() + ) continue if "P3009" in stderr: @@ -818,6 +863,7 @@ class ProxyExtrasDBManager: f"intervention may be required.\n\n" f"Detail: {resolve_err}" ) from resolve_err + budget = budget.after_recovery(f"resolved:{name}") continue if migration_match: migration_name = migration_match.group(1) @@ -831,6 +877,7 @@ class ProxyExtrasDBManager: migration_name, ) ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name) + budget = budget.spend() time.sleep(random.randrange(5, 15)) continue raise RuntimeError( @@ -876,6 +923,7 @@ class ProxyExtrasDBManager: f"intervention may be required.\n\n" f"Detail: {resolve_err}" ) from resolve_err + budget = budget.after_recovery(f"resolved:{name}") continue if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr: @@ -888,6 +936,7 @@ class ProxyExtrasDBManager: ProxyExtrasDBManager._roll_back_migration_best_effort( migration_match.group(1) ) + budget = budget.spend() time.sleep(random.randrange(5, 15)) continue @@ -900,8 +949,9 @@ class ProxyExtrasDBManager: logger.info( "prisma migrate deploy attempt %s deadlocked against " "a concurrent migrate deploy, retrying", - attempt + 1, + budget.attempt_number, ) + budget = budget.spend() time.sleep(random.randrange(5, 15)) continue @@ -909,8 +959,9 @@ class ProxyExtrasDBManager: logger.info( "prisma migrate deploy attempt %s timed out waiting for " "the advisory lock a concurrent migrate deploy holds, retrying", - attempt + 1, + budget.attempt_number, ) + budget = budget.spend() time.sleep(random.randrange(5, 15)) continue @@ -920,9 +971,9 @@ class ProxyExtrasDBManager: ) from e raise RuntimeError( - "Database migration failed after 4 attempts (retry loop " - "exhausted by timeouts, deadlock retries, or repeated " - "idempotent-recovery continues). Check database connectivity, " + f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} " + "attempts that made no progress (timeouts, deadlock retries, or a " + "recovery that had already run once). Check database connectivity, " "load, and _prisma_migrations ledger state, and raise " f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out." ) diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index b3d457707b8..22558faedad 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -703,3 +703,170 @@ class TestSpendLogsPartitionDetectionMissingPsycopg: assert any( "psycopg is not installed" in record.message for record in caplog.records ) + + +_ATTEMPT_BUDGET = 4 + +_P3005_STDERR = """Error: P3005 + +The database schema is not empty. Read more about how to baseline an existing production database: https://pris.ly/d/migrate-baseline +""" + + +def _p3018_stderr(migration_name): + return f"""Error: P3018 + +A migration failed to apply. New migrations cannot be applied before the error is recovered from. + +Migration name: {migration_name} + +Database error code: 42P07 + +Database error: +ERROR: relation "SomeTable" already exists +""" + + +class _MigrateDeployHarness: + """Drives _setup_database_v2 with a scripted sequence of + `prisma migrate deploy` outcomes, with every recovery command faked out so + nothing touches a database or the packaged migrations directory.""" + + def __init__(self, monkeypatch, tmp_path, outcomes, repeat_last=False): + import subprocess as subprocess_module + + import litellm_proxy_extras.utils as utils_module + + self.deploy_calls = [] + self.resolved = [] + self.baselines = 0 + self._outcomes = list(outcomes) + self._repeat_last = repeat_last + self._subprocess_module = subprocess_module + + monkeypatch.delenv("DATABASE_URL", raising=False) + monkeypatch.setattr( + ProxyExtrasDBManager, "_get_prisma_dir", staticmethod(lambda: str(tmp_path)) + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_create_baseline_migration", + staticmethod(self._fake_baseline), + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_roll_back_migration", + staticmethod(lambda name: None), + ) + monkeypatch.setattr( + ProxyExtrasDBManager, + "_resolve_specific_migration", + staticmethod(self.resolved.append), + ) + monkeypatch.setattr(utils_module.subprocess, "run", self._fake_run) + monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None) + + self.baseline_succeeds = True + + def _fake_baseline(self, *args, **kwargs): + self.baselines += 1 + return self.baseline_succeeds + + def _next_outcome(self): + if self._outcomes: + if self._repeat_last and len(self._outcomes) == 1: + return self._outcomes[0] + return self._outcomes.pop(0) + raise AssertionError("prisma migrate deploy called more times than scripted") + + def _fake_run(self, cmd, **kwargs): + assert cmd[1:] == ["migrate", "deploy"], f"unexpected prisma command: {cmd}" + self.deploy_calls.append(cmd) + outcome = self._next_outcome() + if outcome == "ok": + return _FakeCompleted() + if outcome == "timeout": + raise self._subprocess_module.TimeoutExpired(cmd, 1) + raise self._subprocess_module.CalledProcessError(1, cmd, stderr=outcome) + + def run(self): + return ProxyExtrasDBManager._setup_database_v2(use_migrate=True) + + +class TestMigrateDeployAttemptAccounting: + """A `prisma db push` database has a full schema and no ledger, so the v2 + resolver baselines it and then works through every migration whose objects + already exist. Those recoveries make progress, so they must not spend the + retry budget, which is there to stop a run that is getting nowhere.""" + + def test_a_push_created_database_finishes_bootstrapping( + self, monkeypatch, tmp_path + ): + already_there = [ + "20250329084805_new_cron_job_table", + "20250806095134_rename_alias_to_server_name_mcp_table", + "20260224203854_add_agent_object_permissions_table", + "20260301120000_fourth_table", + "20260302120000_fifth_table", + "20260303120000_sixth_table", + ] + harness = _MigrateDeployHarness( + monkeypatch, + tmp_path, + [_P3005_STDERR] + + [_p3018_stderr(name) for name in already_there] + + ["ok"], + ) + + assert harness.run() is True + assert harness.baselines == 1 + assert harness.resolved == already_there + assert len(harness.deploy_calls) == len(already_there) + 2 + + def test_repeated_recovery_of_one_migration_still_gives_up( + self, monkeypatch, tmp_path + ): + harness = _MigrateDeployHarness( + monkeypatch, + tmp_path, + [_p3018_stderr("20250329084805_new_cron_job_table")], + repeat_last=True, + ) + + with pytest.raises(RuntimeError): + harness.run() + assert len(harness.deploy_calls) <= _ATTEMPT_BUDGET + 1 + + def test_timeouts_still_spend_the_budget(self, monkeypatch, tmp_path): + harness = _MigrateDeployHarness( + monkeypatch, tmp_path, ["timeout"], repeat_last=True + ) + + with pytest.raises(RuntimeError): + harness.run() + assert len(harness.deploy_calls) == _ATTEMPT_BUDGET + + def test_a_baseline_that_never_lands_stops_after_the_budget( + self, monkeypatch, tmp_path + ): + harness = _MigrateDeployHarness( + monkeypatch, tmp_path, [_P3005_STDERR], repeat_last=True + ) + harness.baseline_succeeds = False + + with pytest.raises(RuntimeError): + harness.run() + assert len(harness.deploy_calls) == _ATTEMPT_BUDGET + + def test_an_unrecoverable_error_is_not_retried(self, monkeypatch, tmp_path): + harness = _MigrateDeployHarness( + monkeypatch, + tmp_path, + ["Error: P3018\n\nMigration name: 20260101000000_x\n\nERROR: syntax error at or near \"SLECT\"\n"], + repeat_last=True, + ) + + with pytest.raises(RuntimeError): + harness.run() + assert len(harness.deploy_calls) == 1 + assert harness.resolved == [] From 8572544b44662f00f2cea9a573bc6bf9977faa9d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 00:20:46 -0700 Subject: [PATCH 07/26] refactor(proxy-extras): pull the migrate deploy recovery branches into a budget helper _setup_database_v2 decided the next attempt budget inline in eight branches, each rebinding budget before continuing. The branches now live in _budget_after_deploy_failure, which returns the budget the next pass runs under, and the two identical idempotent-recovery blocks share _mark_migration_applied. The loop backs off whenever a pass spent an attempt, which is the same set of paths that slept before. --- .../litellm_proxy_extras/utils.py | 289 ++++++++---------- 1 file changed, 131 insertions(+), 158 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 21f255b0666..a4ea4789b49 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -808,167 +808,16 @@ class ProxyExtrasDBManager: deploy_timeout, PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, ) - budget = budget.spend() - time.sleep(random.randrange(5, 15)) - continue + next_budget = budget.spend() except subprocess.CalledProcessError as e: - stderr = e.stderr or "" + next_budget = ProxyExtrasDBManager._budget_after_deploy_failure( + e, budget, schema_path + ) - if "P3005" in stderr and "database schema is not empty" in stderr: - logger.info( - "Schema exists but no migrations ledger — creating baseline" - ) - baselined = ProxyExtrasDBManager._create_baseline_migration( - schema_path - ) - budget = ( - budget.after_recovery("baseline") - if baselined - else budget.spend() - ) - continue - - if "P3009" in stderr: - migration_match = re.search(r"`(\d+_\S+?)`", stderr) - if ( - migration_match - and ProxyExtrasDBManager._is_idempotent_error(stderr) - ): - name = migration_match.group(1) - logger.info( - f"Migration {name} failed idempotently — marking applied and retrying" - ) - try: - ProxyExtrasDBManager._roll_back_migration(name) - except ( - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - ): - pass # may already be rolled-back - try: - ProxyExtrasDBManager._resolve_specific_migration(name) - except ( - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - ) as resolve_err: - # We're already inside the outer - # `except CalledProcessError` handler — - # re-raising CalledProcessError from here - # would escape as itself, bypassing - # proxy_cli.py's `except RuntimeError`. - raise RuntimeError( - f"Failed to mark migration {name} as applied " - f"after idempotent recovery. Manual " - f"intervention may be required.\n\n" - f"Detail: {resolve_err}" - ) from resolve_err - budget = budget.after_recovery(f"resolved:{name}") - continue - if migration_match: - migration_name = migration_match.group(1) - ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name) - if ledger_logs is not None and ( - ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs - ): - logger.info( - "Migration %s failed in a concurrent migrate deploy " - "deadlock race, rolling its ledger row back and retrying", - migration_name, - ) - ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name) - budget = budget.spend() - time.sleep(random.randrange(5, 15)) - continue - raise RuntimeError( - "Database migration failed and cannot be auto-recovered. " - f"Manual intervention required.\n\nPrisma error:\n{stderr}" - ) from e - - if "P3018" in stderr: - if ProxyExtrasDBManager._is_permission_error(stderr): - raise RuntimeError( - "Database migration failed due to insufficient " - "permissions. Please grant the required privileges " - f"and retry.\n\nPrisma error:\n{stderr}" - ) from e - - migration_match = re.search( - r"Migration name: (\d+_\S+)", stderr - ) - if ( - migration_match - and ProxyExtrasDBManager._is_idempotent_error(stderr) - ): - name = migration_match.group(1) - logger.info( - f"Migration {name} SQL hit idempotent error — marking applied and retrying" - ) - try: - ProxyExtrasDBManager._roll_back_migration(name) - except ( - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - ): - pass # may already be rolled-back - try: - ProxyExtrasDBManager._resolve_specific_migration(name) - except ( - subprocess.CalledProcessError, - subprocess.TimeoutExpired, - ) as resolve_err: - raise RuntimeError( - f"Failed to mark migration {name} as applied " - f"after idempotent recovery. Manual " - f"intervention may be required.\n\n" - f"Detail: {resolve_err}" - ) from resolve_err - budget = budget.after_recovery(f"resolved:{name}") - continue - - if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr: - logger.info( - "Migration %s deadlocked against a concurrent " - "migrate deploy, rolling its ledger row back " - "and retrying", - migration_match.group(1), - ) - ProxyExtrasDBManager._roll_back_migration_best_effort( - migration_match.group(1) - ) - budget = budget.spend() - time.sleep(random.randrange(5, 15)) - continue - - raise RuntimeError( - "Database migration failed and cannot be auto-recovered. " - f"Manual intervention required.\n\nPrisma error:\n{stderr}" - ) from e - - if _MIGRATION_DEADLOCK_MARKER in stderr: - logger.info( - "prisma migrate deploy attempt %s deadlocked against " - "a concurrent migrate deploy, retrying", - budget.attempt_number, - ) - budget = budget.spend() - time.sleep(random.randrange(5, 15)) - continue - - if "P1002" in stderr and "advisory lock" in stderr: - logger.info( - "prisma migrate deploy attempt %s timed out waiting for " - "the advisory lock a concurrent migrate deploy holds, retrying", - budget.attempt_number, - ) - budget = budget.spend() - time.sleep(random.randrange(5, 15)) - continue - - raise RuntimeError( - "Database migration failed and cannot be auto-recovered. " - f"Manual intervention required.\n\nPrisma error:\n{stderr}" - ) from e + if next_budget.attempts_left < budget.attempts_left: + time.sleep(random.randrange(5, 15)) + budget = next_budget # rebind-ok: the loop carries the budget from one migrate deploy pass to the next raise RuntimeError( f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} " @@ -980,6 +829,130 @@ class ProxyExtrasDBManager: finally: os.chdir(original_dir) + @staticmethod + def _budget_after_deploy_failure( + error: subprocess.CalledProcessError, + budget: "_MigrateAttemptBudget", + schema_path: str, + ) -> "_MigrateAttemptBudget": + """Recover from one failed `prisma migrate deploy`, and price the pass. + + Returns the budget the next pass runs under, or raises when the failure + is not one this resolver knows how to recover from. + """ + stderr = error.stderr or "" + + if "P3005" in stderr and "database schema is not empty" in stderr: + logger.info("Schema exists but no migrations ledger — creating baseline") + if ProxyExtrasDBManager._create_baseline_migration(schema_path): + return budget.after_recovery("baseline") + return budget.spend() + + if "P3009" in stderr: + migration_match = re.search(r"`(\d+_\S+?)`", stderr) + if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr): + name = migration_match.group(1) + logger.info( + f"Migration {name} failed idempotently — marking applied and retrying" + ) + ProxyExtrasDBManager._mark_migration_applied(name) + return budget.after_recovery(f"resolved:{name}") + if migration_match: + migration_name = migration_match.group(1) + ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name) + if ledger_logs is not None and ( + ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs + ): + logger.info( + "Migration %s failed in a concurrent migrate deploy " + "deadlock race, rolling its ledger row back and retrying", + migration_name, + ) + ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name) + return budget.spend() + raise RuntimeError( + "Database migration failed and cannot be auto-recovered. " + f"Manual intervention required.\n\nPrisma error:\n{stderr}" + ) from error + + if "P3018" in stderr: + if ProxyExtrasDBManager._is_permission_error(stderr): + raise RuntimeError( + "Database migration failed due to insufficient " + "permissions. Please grant the required privileges " + f"and retry.\n\nPrisma error:\n{stderr}" + ) from error + + migration_match = re.search(r"Migration name: (\d+_\S+)", stderr) + if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr): + name = migration_match.group(1) + logger.info( + f"Migration {name} SQL hit idempotent error — marking applied and retrying" + ) + ProxyExtrasDBManager._mark_migration_applied(name) + return budget.after_recovery(f"resolved:{name}") + + if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr: + logger.info( + "Migration %s deadlocked against a concurrent " + "migrate deploy, rolling its ledger row back " + "and retrying", + migration_match.group(1), + ) + ProxyExtrasDBManager._roll_back_migration_best_effort( + migration_match.group(1) + ) + return budget.spend() + + raise RuntimeError( + "Database migration failed and cannot be auto-recovered. " + f"Manual intervention required.\n\nPrisma error:\n{stderr}" + ) from error + + if _MIGRATION_DEADLOCK_MARKER in stderr: + logger.info( + "prisma migrate deploy attempt %s deadlocked against " + "a concurrent migrate deploy, retrying", + budget.attempt_number, + ) + return budget.spend() + + if "P1002" in stderr and "advisory lock" in stderr: + logger.info( + "prisma migrate deploy attempt %s timed out waiting for " + "the advisory lock a concurrent migrate deploy holds, retrying", + budget.attempt_number, + ) + return budget.spend() + + raise RuntimeError( + "Database migration failed and cannot be auto-recovered. " + f"Manual intervention required.\n\nPrisma error:\n{stderr}" + ) from error + + @staticmethod + def _mark_migration_applied(name: str) -> None: + """Roll a failed ledger row back if it is still there, then mark it applied.""" + try: + ProxyExtrasDBManager._roll_back_migration(name) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired): + pass # may already be rolled-back + try: + ProxyExtrasDBManager._resolve_specific_migration(name) + except ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + ) as resolve_err: + # We're called from inside an `except CalledProcessError` handler — + # re-raising CalledProcessError from here would escape as itself, + # bypassing proxy_cli.py's `except RuntimeError`. + raise RuntimeError( + f"Failed to mark migration {name} as applied " + f"after idempotent recovery. Manual " + f"intervention may be required.\n\n" + f"Detail: {resolve_err}" + ) from resolve_err + @staticmethod def apply_replica_identity_full_if_requested() -> bool: """ From 658f50663d19f613a3f5caf998168da019764ad8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 3 Sep 2026 00:35:32 -0700 Subject: [PATCH 08/26] fix(ui): keep Virtual Keys list state in the URL so it survives leaving the page (#39481) * fix(ui): keep Virtual Keys list state in the URL so it survives leaving the page The search term, sort, pagination and drawer filters lived in component state, so navigating away from Virtual Keys and back reset the table to an unfiltered first page. Move them into query state alongside the existing ?key= deep link, which also makes a filtered view shareable. * fix(ui): namespace the Virtual Keys filter params and bound page inputs The unprefixed team_id filter hijacked the /api-keys create-key deep link, which already takes team_id as a prefill, so ?create=true&team_id=X silently filtered the list underneath the modal. Prefix the four drawer filters. Now that page and page_size come from the address bar, clamp them to what /key/list accepts instead of forwarding 0, negatives or an int64-overflowing page straight through, and trim filter values arriving from a URL the same way the drawer already trims them. * fix(ui): fall back to a sortable column when the URL names an unknown one A hand-edited or stale sort_by reached /key/list, which 400s it, leaving the Virtual Keys page on its loading skeleton with no error. Validate it against the fields the table's own headers can produce, and clear sort_by rather than blanking it when a sort is reset so the URL stays clean. Also replaces a default-state URL assertion that ran before any query-state write could land, so it could not fail for the regression it named. * fix(ui): use TanStack's functionalUpdate instead of a hand-rolled updater resolver The local helper narrowed typeof updater === "function" against an unconstrained T, which TypeScript cannot do because T itself may be a function type, so next build failed to type check. table-core already exports the same helper. --- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 196 +++++++++++++++++- .../VirtualKeysPage/VirtualKeysTable.tsx | 158 ++++++++++---- .../VirtualKeysPage/keyTableColumns.tsx | 8 + 3 files changed, 317 insertions(+), 45 deletions(-) diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index b06448912d0..93f860333d8 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -4,6 +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 { 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"; @@ -176,8 +177,10 @@ const keysResult = (keys: KeyResponse[], data: Partial = {}, extra const openFilters = () => fireEvent.click(screen.getByRole("button", { name: "Filters" })); -const lastKeyParam = (onUrlUpdate: Mock) => - onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("key"); +const lastSearchParam = (onUrlUpdate: Mock, name: string) => + onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get(name); + +const lastKeyParam = (onUrlUpdate: Mock) => lastSearchParam(onUrlUpdate, "key"); beforeEach(() => { vi.clearAllMocks(); @@ -510,7 +513,8 @@ describe("server-side filtering – the LIT-4080 regression guard", () => { }); it("drops the filter from the useKeys query when it is cleared", async () => { - renderWithProviders(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); openFilters(); const userIdInput = await screen.findByPlaceholderText(/Enter User ID/); @@ -520,6 +524,12 @@ describe("server-side filtering – the LIT-4080 regression guard", () => { await waitFor(() => { expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ userID: "user-42" })); }); + // Let the filter reach the URL before clearing it: NuqsTestingAdapter runs + // resetUrlUpdateQueueOnMount on every render, so a still-queued write can be + // aborted by the re-render its own predecessor triggers. + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_user")).toBe("user-42"); + }); fireEvent.click(screen.getByTestId("datatable-clear-filters")); @@ -638,3 +648,183 @@ describe("Status column reflects blocked / expiry / scim metadata", () => { }); }); }); + +describe("table state lives in the URL so it survives leaving and returning to the page", () => { + it("restores the search term, sort and pagination from the URL on mount", async () => { + renderWithProviders(, { + searchParams: { key_search: "prod", sort_by: "spend", sort_order: "asc", page: "3", page_size: "25" }, + }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 3, + 25, + expect.objectContaining({ selectedKeyAlias: "prod", sortBy: "spend", sortOrder: "asc" }), + ); + }); + expect(screen.getByPlaceholderText(/Search by key alias/)).toHaveValue("prod"); + }); + + it("restores the drawer filters from the URL on mount", async () => { + renderWithProviders(, { searchParams: { filter_team: "team-1", filter_user: "user-42" } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ teamID: "team-1", userID: "user-42" }), + ); + }); + expect(screen.getByTestId("filter-chip-team_id")).toHaveTextContent("Test Team"); + }); + + it("writes the search term to the URL", async () => { + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + + fireEvent.change(screen.getByPlaceholderText(/Search by key alias/), { target: { value: "prod" } }); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "key_search")).toBe("prod"); + }); + }); + + it("writes the sort field and direction to the URL", async () => { + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + + fireEvent.click(screen.getByRole("button", { name: "Key" })); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "sort_by")).toBe("key_alias"); + }); + expect(lastSearchParam(onUrlUpdate, "sort_order")).toBe("asc"); + }); + + it("writes an applied drawer filter to the URL and clears it again", async () => { + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + + openFilters(); + fireEvent.change(await screen.findByPlaceholderText(/Enter User ID/), { target: { value: "user-42" } }); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_user")).toBe("user-42"); + }); + + fireEvent.click(screen.getByTestId("datatable-clear-filters")); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_user")).toBeNull(); + }); + expect(screen.queryByTestId("filter-chip-user_id")).not.toBeInTheDocument(); + }); + + it("returns to page 1 when the search term changes", async () => { + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: { page: "3" }, onUrlUpdate }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(3, 50, expect.anything()); + }); + + fireEvent.change(screen.getByPlaceholderText(/Search by key alias/), { target: { value: "prod" } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ selectedKeyAlias: "prod" })); + }); + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "page")).toBeNull(); + }); + }); + + it("leaves the create-key deep link's team_id alone instead of filtering the list with it", async () => { + renderWithProviders(, { searchParams: { create: "true", team_id: "team-1" } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ teamID: undefined })); + }); + expect(screen.queryByTestId("filter-chip-team_id")).not.toBeInTheDocument(); + }); + + it.each([ + ["0", 1], + ["-3", 1], + ])("clamps a hand-edited page of %s up to the first page", async (page, expected) => { + renderWithProviders(, { searchParams: { page } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(expected, 50, expect.anything()); + }); + }); + + it.each([ + ["0", 1], + ["1000", 100], + ])("clamps a hand-edited page_size of %s into the range /key/list accepts", async (pageSize, expected) => { + renderWithProviders(, { searchParams: { page_size: pageSize } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, expected, expect.anything()); + }); + }); + + it("trims whitespace off a filter that arrived from the URL", async () => { + renderWithProviders(, { searchParams: { filter_user: " user-42 " } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ userID: "user-42" })); + }); + }); + + 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" } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith( + 1, + 50, + expect.objectContaining({ sortBy: "created_at", sortOrder: "desc" }), + ); + }); + expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); + }); + + it.each(KEY_TABLE_SORT_FIELDS)("round-trips a %s sort from the URL", async (field) => { + renderWithProviders(, { searchParams: { sort_by: field, sort_order: "asc" } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: field, sortOrder: "asc" })); + }); + }); + + it("clears sort_by from the URL when the Spend / Budget sort is reset", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Spend ascending", "menuitem"); + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "sort_by")).toBe("spend"); + }); + + await chooseSelectOption(user, screen.getByTestId("sort-trigger-spend"), "Reset", "menuitem"); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "sort_by")).toBeNull(); + }); + expect(lastSearchParam(onUrlUpdate, "sort_order")).toBeNull(); + }); + + it("drops the search param back out of the URL when the search box is cleared", async () => { + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: { key_search: "prod" }, onUrlUpdate }); + + fireEvent.change(screen.getByPlaceholderText(/Search by key alias/), { target: { value: "" } }); + + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "key_search")).toBeNull(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 7dd25ca1fd0..8ec8b6d0c3f 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -15,34 +15,65 @@ import { SearchSelect } from "@/components/shared/SearchSelect"; import { PageHeader } from "@/components/shared/PageHeader"; import { Input } from "@/components/ui/input"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; -import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; +import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; import { KeyRound } from "lucide-react"; -import { parseAsString, useQueryState } from "nuqs"; +import { createParser, parseAsInteger, parseAsString, parseAsStringLiteral, useQueryState, useQueryStates } from "nuqs"; import React, { useCallback, useMemo, useState } from "react"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; import KeyInfoView from "../templates/key_info_view"; -import { getKeyTableColumns, KEY_TABLE_HIDDEN_COLUMNS } from "./keyTableColumns"; +import { getKeyTableColumns, KEY_TABLE_HIDDEN_COLUMNS, KEY_TABLE_SORT_FIELDS } from "./keyTableColumns"; interface VirtualKeysTableProps { headerActions?: React.ReactNode; } -const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; +const FILTER_COLUMNS = ["team_id", "org_id", "user_id", "key_hash"] as const; +type FilterColumn = (typeof FILTER_COLUMNS)[number]; -const toSortOrder = (sorting: SortingState): "asc" | "desc" | undefined => { - const active = sorting[0]; - if (!active) return undefined; - return active.desc ? "desc" : "asc"; -}; - -const FILTER_LABELS: Record = { +const FILTER_LABELS: Record = { team_id: "Team", org_id: "Organization", user_id: "User ID", key_hash: "Key ID", }; +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 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(""), +}; + +const toSortOrder = (active: SortingState[number]): "asc" | "desc" => (active.desc ? "desc" : "asc"); + +const filterValue = (filters: ColumnFiltersState, column: FilterColumn): string | null => { + const value = filters.find((filter) => filter.id === column)?.value; + return (typeof value === "string" ? value.trim() : "") || null; +}; + export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const { data: fetchedOrganizations } = useOrganizations(); const organizations = useMemo(() => fetchedOrganizations ?? [], [fetchedOrganizations]); @@ -50,32 +81,48 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const allTeams = useMemo(() => fetchedTeams ?? [], [fetchedTeams]); const [selectedKeyId, setSelectedKeyId] = useQueryState("key", parseAsString.withOptions({ history: "push" })); - const [sorting, setSorting] = useState(DEFAULT_SORTING); - const [tablePagination, setTablePagination] = useState({ pageIndex: 0, pageSize: 50 }); - const [columnFilters, setColumnFilters] = useState([]); + const [tableState, setTableState] = useQueryStates(TABLE_STATE); const [filtersOpen, setFiltersOpen] = useState(false); - const [searchInput, setSearchInput] = useState(""); + const searchInput = tableState.key_search; const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); - const getFilterValue = useCallback( - (columnId: string): string | undefined => { - const entry = columnFilters.find((filter) => filter.id === columnId); - return typeof entry?.value === "string" && entry.value.trim() ? entry.value.trim() : undefined; - }, - [columnFilters], + // 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 } = 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(), + }), + [filter_team, filter_org, filter_user, filter_key_id], + ); + const columnFilters = useMemo( + () => + FILTER_COLUMNS.filter((column) => appliedFilters[column]).map((column) => ({ + id: column, + value: appliedFilters[column], + })), + [appliedFilters], ); - const sortBy = sorting[0]?.id; - const sortOrder = toSortOrder(sorting); - const keyListOptions = { - teamID: getFilterValue("team_id"), - organizationID: getFilterValue("org_id"), + teamID: appliedFilters.team_id || undefined, + organizationID: appliedFilters.org_id || undefined, selectedKeyAlias: searchQuery.trim() || undefined, - userID: getFilterValue("user_id"), - keyHash: getFilterValue("key_hash"), + userID: appliedFilters.user_id || undefined, + keyHash: appliedFilters.key_hash || undefined, sortBy, - sortOrder, + sortOrder: tableState.sort_order, expand: "user", }; @@ -89,20 +136,47 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const keyList = useMemo(() => keys?.keys ?? [], [keys]); const rowCount = keys?.total_count ?? 0; - const handleSearchChange = useCallback((value: string) => { - setSearchInput(value); - setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); - }, []); + const handleSearchChange = useCallback( + (value: string) => { + void setTableState({ key_search: value || null, page: null }); + }, + [setTableState], + ); - const handleSortingChange = useCallback>((updaterOrValue) => { - setSorting(updaterOrValue); - setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); - }, []); + 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) => { - setColumnFilters(updaterOrValue); - setTablePagination((prev) => ({ ...prev, pageIndex: 0 })); - }, []); + 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"), + 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) }), @@ -199,7 +273,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { onSortingChange={handleSortingChange} paginationMode="server" pagination={tablePagination} - onPaginationChange={setTablePagination} + onPaginationChange={handlePaginationChange} rowCount={rowCount} filterMode="server" columnFilters={columnFilters} diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index 51b8b734b4c..0865fe76519 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -32,6 +32,14 @@ const SPEND_BUDGET_SORT_FIELDS: DataTableSortField[] = [ { id: "max_budget", label: "Budget" }, ]; +export const KEY_TABLE_SORT_FIELDS: readonly string[] = [ + "key_alias", + "token", + "created_at", + "updated_at", + ...SPEND_BUDGET_SORT_FIELDS.map((field) => field.id), +]; + const getKeyStatus = (key: KeyResponse): KeyStatus => { if (key.blocked === true) { const isScimBlocked = (key.metadata as Record | null | undefined)?.scim_blocked === true; From ccbd3e495c203484acd0b96ed296d82377796b7d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:24:09 -0700 Subject: [PATCH 09/26] fix(agents): keep the published agent in public_agent_groups `POST /v1/agents/{id}/make_public` appended the agent id to `litellm.public_agent_groups` and only then called `get_config()`, which re-applies the DB's `litellm_settings` over the module globals and threw the append away. The config it saved was therefore a no-op: the endpoint answered 200 with an empty `public_agent_groups`, the agent never reached `GET /public/agent_hub`, and re-publishing never hit the "already public" 400. Read the config first, derive the new list from the refreshed globals, save it, then update the global Also fixes the e2e model hub spec, which is flaky for a second reason: the "Make Models Public" modal preselects the groups that are already public, so a blind click on "Select All" cleared them and left "Next" disabled for the full 15s action timeout. Check the box instead of toggling it, and wait for "Next" to be enabled before clicking --- litellm/proxy/agent_endpoints/endpoints.py | 25 ++++--- tests/e2e/ui/tests/modelHub/modelHub.spec.ts | 9 ++- .../proxy/agent_endpoints/test_endpoints.py | 69 +++++++++++++++++++ 3 files changed, 87 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 3b8151d1064..ec50500fdb2 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -939,35 +939,34 @@ async def make_agent_public( if agent is None: raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") - if litellm.public_agent_groups is None: - litellm.public_agent_groups = [] - # handle duplicates - if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups): + # get_config() re-applies the DB's litellm_settings over the in-memory + # globals, so read it before deriving the new list and assign the global after + config: Final = await proxy_config.get_config() + + current_public_agent_groups: Final = list(litellm.public_agent_groups or []) + if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(current_public_agent_groups): raise HTTPException( status_code=400, detail=f"Agent with name {agent.agent_name} already in public agent groups", ) - litellm.public_agent_groups.append(agent.agent_id) + updated_public_agent_groups: Final = [*current_public_agent_groups, agent.agent_id] - # Load existing config - config: Final = await proxy_config.get_config() - - # Update config with new settings if "litellm_settings" not in config or config["litellm_settings"] is None: config["litellm_settings"] = {} - config["litellm_settings"]["public_agent_groups"] = litellm.public_agent_groups + config["litellm_settings"]["public_agent_groups"] = updated_public_agent_groups - # Save the updated config await proxy_config.save_config(new_config=config) + litellm.public_agent_groups = updated_public_agent_groups + verbose_proxy_logger.debug( - "Updated public agent groups to: %s by user: %s", litellm.public_agent_groups, user_api_key_dict.user_id + "Updated public agent groups to: %s by user: %s", updated_public_agent_groups, user_api_key_dict.user_id ) return { "message": "Successfully updated public agent groups", - "public_agent_groups": litellm.public_agent_groups, + "public_agent_groups": updated_public_agent_groups, "updated_by": user_api_key_dict.user_id, } except HTTPException: diff --git a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts index 6877fc9c48d..7c43b9f761c 100644 --- a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts +++ b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts @@ -22,11 +22,14 @@ test.describe("AI Hub (internal admin view)", () => { // on the disabled-Next button or the success toast. await expect(modal.getByText(/Select All \(\d+\)/)).toBeVisible({ timeout: 5_000 }); - // Step 1: pick the seeded models via "Select All" - await modal.getByText(/Select All/i).click(); + // Step 1: pick the seeded models via "Select All". check() rather than click() because + // the modal preselects groups that are already public, and a click would clear them. + await modal.getByRole("checkbox", { name: /Select All/ }).check(); // Move to confirm step - await modal.getByRole("button", { name: "Next" }).click(); + const next = modal.getByRole("button", { name: "Next" }); + await expect(next).toBeEnabled(); + await next.click(); await expect(modal.getByText("Confirm Making Models Public")).toBeVisible({ timeout: 5_000 }); // Submit diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index a78b3238a9a..7c1db851a48 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -1062,3 +1062,72 @@ def test_merged_agent_card_url_has_no_double_slash_without_proxy_base_url( interface_url = merged["supportedInterfaces"][0]["url"] assert interface_url == f"{base_url.rstrip('/')}/a2a/agent-xyz" assert "//a2a" not in interface_url + + +class _DbBackedProxyConfig: + """Round-trips `litellm_settings` through the DB overlay the proxy applies on every + `get_config()`, which is what re-assigns the `litellm.public_*` globals in production.""" + + def __init__(self) -> None: + self.stored_litellm_settings: dict = {} + + async def get_config(self) -> dict: + from litellm.proxy.proxy_server import ProxyConfig + + config: dict = {"litellm_settings": {}} + if not self.stored_litellm_settings: + return config + return ProxyConfig()._update_config_fields( + current_config=config, + param_name="litellm_settings", + db_param_value=dict(self.stored_litellm_settings), + ) + + async def save_config(self, new_config: dict) -> None: + self.stored_litellm_settings = dict(new_config.get("litellm_settings") or {}) + + +def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch): + """A second /make_public call must not drop the agent published by the first one.""" + import litellm + from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module + from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry + + registry = AgentRegistry() + registry.register_agent(_sample_agent_response(agent_id="agent-1", agent_name="Agent One")) + registry.register_agent(_sample_agent_response(agent_id="agent-2", agent_name="Agent Two")) + + monkeypatch.setattr(agent_registry_module, "global_agent_registry", registry) + monkeypatch.setattr(litellm, "public_agent_groups", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", _DbBackedProxyConfig()) + + first = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) + second = client.post("/v1/agents/agent-2/make_public", headers={"Authorization": "Bearer test-key"}) + + assert first.status_code == 200 + assert second.status_code == 200 + assert second.json()["public_agent_groups"] == ["agent-1", "agent-2"] + assert [agent.agent_id for agent in registry.get_public_agent_list()] == ["agent-1", "agent-2"] + + +def test_make_agent_public_rejects_an_already_public_agent(monkeypatch): + """The duplicate guard must still fire when the published list comes back from the DB.""" + import litellm + from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module + from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry + + registry = AgentRegistry() + registry.register_agent(_sample_agent_response(agent_id="agent-1", agent_name="Agent One")) + + monkeypatch.setattr(agent_registry_module, "global_agent_registry", registry) + monkeypatch.setattr(litellm, "public_agent_groups", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", _DbBackedProxyConfig()) + + first = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) + duplicate = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) + + assert first.status_code == 200 + assert duplicate.status_code == 400 + assert "already in public agent groups" in duplicate.json()["detail"] From 9d862a65831cf1a7dc8ba6a7e69900d3774dc85e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:27:43 -0700 Subject: [PATCH 10/26] style: drop explanatory comments from the agent publish fix --- litellm/proxy/agent_endpoints/endpoints.py | 2 -- tests/e2e/ui/tests/modelHub/modelHub.spec.ts | 3 +-- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index ec50500fdb2..cc17672553b 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -939,8 +939,6 @@ async def make_agent_public( if agent is None: raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") - # get_config() re-applies the DB's litellm_settings over the in-memory - # globals, so read it before deriving the new list and assign the global after config: Final = await proxy_config.get_config() current_public_agent_groups: Final = list(litellm.public_agent_groups or []) diff --git a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts index 7c43b9f761c..1fa3e4c530e 100644 --- a/tests/e2e/ui/tests/modelHub/modelHub.spec.ts +++ b/tests/e2e/ui/tests/modelHub/modelHub.spec.ts @@ -22,8 +22,7 @@ test.describe("AI Hub (internal admin view)", () => { // on the disabled-Next button or the success toast. await expect(modal.getByText(/Select All \(\d+\)/)).toBeVisible({ timeout: 5_000 }); - // Step 1: pick the seeded models via "Select All". check() rather than click() because - // the modal preselects groups that are already public, and a click would clear them. + // Step 1: pick the seeded models via "Select All" await modal.getByRole("checkbox", { name: /Select All/ }).check(); // Move to confirm step From b1695e909070562dffb6d01655676a2fa0b8743a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:33:55 -0700 Subject: [PATCH 11/26] test(agents): type the public-agent regression tests fully --- .../proxy/agent_endpoints/test_endpoints.py | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 7c1db851a48..51797cbb308 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -1,3 +1,4 @@ +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1069,12 +1070,12 @@ class _DbBackedProxyConfig: `get_config()`, which is what re-assigns the `litellm.public_*` globals in production.""" def __init__(self) -> None: - self.stored_litellm_settings: dict = {} + self.stored_litellm_settings: dict[str, object] = {} - async def get_config(self) -> dict: + async def get_config(self) -> dict[str, dict[str, object]]: from litellm.proxy.proxy_server import ProxyConfig - config: dict = {"litellm_settings": {}} + config: Final[dict[str, dict[str, object]]] = {"litellm_settings": {}} if not self.stored_litellm_settings: return config return ProxyConfig()._update_config_fields( @@ -1083,17 +1084,17 @@ class _DbBackedProxyConfig: db_param_value=dict(self.stored_litellm_settings), ) - async def save_config(self, new_config: dict) -> None: + async def save_config(self, new_config: dict[str, dict[str, object]]) -> None: self.stored_litellm_settings = dict(new_config.get("litellm_settings") or {}) -def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch): +def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch: pytest.MonkeyPatch) -> None: """A second /make_public call must not drop the agent published by the first one.""" import litellm from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry - registry = AgentRegistry() + registry: Final = AgentRegistry() registry.register_agent(_sample_agent_response(agent_id="agent-1", agent_name="Agent One")) registry.register_agent(_sample_agent_response(agent_id="agent-2", agent_name="Agent Two")) @@ -1102,8 +1103,8 @@ def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", _DbBackedProxyConfig()) - first = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) - second = client.post("/v1/agents/agent-2/make_public", headers={"Authorization": "Bearer test-key"}) + first: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) + second: Final = client.post("/v1/agents/agent-2/make_public", headers={"Authorization": "Bearer test-key"}) assert first.status_code == 200 assert second.status_code == 200 @@ -1111,13 +1112,13 @@ def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch): assert [agent.agent_id for agent in registry.get_public_agent_list()] == ["agent-1", "agent-2"] -def test_make_agent_public_rejects_an_already_public_agent(monkeypatch): +def test_make_agent_public_rejects_an_already_public_agent(monkeypatch: pytest.MonkeyPatch) -> None: """The duplicate guard must still fire when the published list comes back from the DB.""" import litellm from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry - registry = AgentRegistry() + registry: Final = AgentRegistry() registry.register_agent(_sample_agent_response(agent_id="agent-1", agent_name="Agent One")) monkeypatch.setattr(agent_registry_module, "global_agent_registry", registry) @@ -1125,8 +1126,8 @@ def test_make_agent_public_rejects_an_already_public_agent(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", _DbBackedProxyConfig()) - first = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) - duplicate = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) + first: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) + duplicate: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) assert first.status_code == 200 assert duplicate.status_code == 400 From 15e956db33829cc82800e36cd89dd23f3d8701b4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 05:21:42 -0700 Subject: [PATCH 12/26] test(agents): make the make_public regression tests fail without the fix The config stub shared one list object between save_config and get_config, so the DB overlay handed the endpoint back the very list it had just appended to and both tests passed with the product fix reverted. Store the settings as JSON the way the litellm_config row does, and check the duplicate guard against a list that only ever existed in the DB. --- .../proxy/agent_endpoints/test_endpoints.py | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 51797cbb308..067f5a9f64c 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -1,3 +1,4 @@ +import json from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -1067,25 +1068,29 @@ def test_merged_agent_card_url_has_no_double_slash_without_proxy_base_url( class _DbBackedProxyConfig: """Round-trips `litellm_settings` through the DB overlay the proxy applies on every - `get_config()`, which is what re-assigns the `litellm.public_*` globals in production.""" + `get_config()`, which is what re-assigns the `litellm.public_*` globals in production. - def __init__(self) -> None: - self.stored_litellm_settings: dict[str, object] = {} + Storage goes through JSON the way the `litellm_config` row does, so every read hands back + freshly built values instead of the objects the endpoint still holds a reference to.""" + + def __init__(self, stored_litellm_settings: dict[str, object] | None = None) -> None: + self.stored_litellm_settings_json: str = json.dumps(stored_litellm_settings or {}) async def get_config(self) -> dict[str, dict[str, object]]: from litellm.proxy.proxy_server import ProxyConfig config: Final[dict[str, dict[str, object]]] = {"litellm_settings": {}} - if not self.stored_litellm_settings: + db_param_value: Final[dict[str, object]] = json.loads(self.stored_litellm_settings_json) + if not db_param_value: return config return ProxyConfig()._update_config_fields( current_config=config, param_name="litellm_settings", - db_param_value=dict(self.stored_litellm_settings), + db_param_value=db_param_value, ) async def save_config(self, new_config: dict[str, dict[str, object]]) -> None: - self.stored_litellm_settings = dict(new_config.get("litellm_settings") or {}) + self.stored_litellm_settings_json = json.dumps(new_config.get("litellm_settings") or {}) def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch: pytest.MonkeyPatch) -> None: @@ -1112,8 +1117,8 @@ def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch: pytest.Mo assert [agent.agent_id for agent in registry.get_public_agent_list()] == ["agent-1", "agent-2"] -def test_make_agent_public_rejects_an_already_public_agent(monkeypatch: pytest.MonkeyPatch) -> None: - """The duplicate guard must still fire when the published list comes back from the DB.""" +def test_make_agent_public_rejects_an_agent_published_only_in_the_db(monkeypatch: pytest.MonkeyPatch) -> None: + """The duplicate guard must fire off the stored list, not just what this process published.""" import litellm from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry @@ -1124,11 +1129,12 @@ def test_make_agent_public_rejects_an_already_public_agent(monkeypatch: pytest.M monkeypatch.setattr(agent_registry_module, "global_agent_registry", registry) monkeypatch.setattr(litellm, "public_agent_groups", None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", _DbBackedProxyConfig()) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config", + _DbBackedProxyConfig({"public_agent_groups": ["agent-1"]}), + ) - first: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) duplicate: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"}) - assert first.status_code == 200 assert duplicate.status_code == 400 assert "already in public agent groups" in duplicate.json()["detail"] From 418d0e79ba9b65a1ba616b2cad4608cba5d0e9b1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 05:30:28 -0700 Subject: [PATCH 13/26] fix(proxy): answer 404 when deleting a credential that was never stored prisma's `delete` returns None when the `where` clause matched no row instead of raising, and the handler never looked at the return value. It went straight on to filter an in-memory list that never held the name and answered 200 "Credential deleted successfully", so an operator scripting credential cleanup could not tell a real deletion from a typo. Look at what the repository returned and answer 404 with the name, the same rejection PATCH /credentials/{credential_name} already gives. A credential that only exists in the config yaml is never written to the table, so it now answers 404 too, which is honest: reporting success for it is the same lie, since it comes back on the next proxy boot. Adds regression tests for the delete 404, the still-working delete, the config-yaml-only credential, and for the raise-not-return fix on both DELETE /credentials/{credential_name} and GET /credentials, which serialized a rejection as the 200 response body. --- .../proxy/credential_endpoints/endpoints.py | 7 +- .../credential_endpoints/test_endpoints.py | 171 ++++++++++++++---- 2 files changed, 146 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 106b43a722c..0b9634b15c9 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -234,7 +234,12 @@ async def delete_credential( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - await CredentialsRepository(prisma_client).delete_by_name(credential_name) + deleted: Final = await CredentialsRepository(prisma_client).delete_by_name(credential_name) + if deleted is None: + raise HTTPException( + status_code=404, + detail="Credential not found. Got credential name: " + credential_name, + ) ## DELETE FROM LITELLM ## litellm.credential_list = [cred for cred in litellm.credential_list if cred.credential_name != credential_name] diff --git a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py index e2fa1de6962..cf89178c766 100644 --- a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py @@ -9,6 +9,7 @@ from fastapi.testclient import TestClient sys.path.insert(0, os.path.abspath("../../../..")) +import litellm from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app @@ -21,16 +22,12 @@ def _as_admin(): return UserAPIKeyAuth(api_key="test-key", user_role="proxy_admin") -def _patch_credential(name: str, body: dict): +def _call_as_admin(method: str, path: str, json_body: dict | None = None): missing = object() previous_override = app.dependency_overrides.get(user_api_key_auth, missing) app.dependency_overrides[user_api_key_auth] = _as_admin try: - return client.patch( - f"/credentials/{name}", - json=body, - headers={"Authorization": "Bearer test-key"}, - ) + return client.request(method, path, json=json_body, headers={"Authorization": "Bearer test-key"}) finally: if previous_override is missing: app.dependency_overrides.pop(user_api_key_auth, None) @@ -38,37 +35,69 @@ def _patch_credential(name: str, body: dict): app.dependency_overrides[user_api_key_auth] = previous_override -def test_update_credential_answers_404_when_the_credential_does_not_exist(): +def _patch_credential(name: str, body: dict): + return _call_as_admin("PATCH", f"/credentials/{name}", body) + + +def _delete_credential(name: str): + return _call_as_admin("DELETE", f"/credentials/{name}") + + +def _list_credentials(): + return _call_as_admin("GET", "/credentials") + + +@pytest.fixture +def credential_store(): + """Stands the credential store up for one test: whether the database is reachable, what + the proxy is already serving from memory, and what each repository call hands back.""" + + def install( + *, + connected: bool = True, + in_memory: tuple[object, ...] = (), + **repository_calls: AsyncMock, + ) -> None: + patch("litellm.proxy.proxy_server.prisma_client", MagicMock() if connected else None).start() + patch("litellm.proxy.proxy_server.master_key", "sk-test-master").start() + patch.object(litellm, "credential_list", list(in_memory)).start() + repository = patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository").start() + for call_name, result in repository_calls.items(): + setattr(repository.return_value, call_name, result) + + yield install + patch.stopall() + + +def test_update_credential_answers_404_when_the_credential_does_not_exist(credential_store): """Regression: the handler used to ``return handle_exception_on_proxy(e)``, which makes the exception the response body and lets FastAPI answer 200, so a write the handler rejected read as a success to every caller that checks the status. The dashboard's API client branches on the status, so it reported a failed edit as applied.""" - with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( - "litellm.proxy.credential_endpoints.endpoints.CredentialsRepository" - ) as repository: - repository.return_value.find_by_name = AsyncMock(return_value=None) + credential_store(find_by_name=AsyncMock(return_value=None)) - response = _patch_credential( - "definitely-not-there", - {"credential_name": "definitely-not-there", "credential_values": {"api_key": "sk-x"}, "credential_info": {}}, - ) + response = _patch_credential( + "definitely-not-there", + {"credential_name": "definitely-not-there", "credential_values": {"api_key": "sk-x"}, "credential_info": {}}, + ) assert response.status_code == 404, f"rejected write answered {response.status_code}: {response.text}" assert "error" in response.json() -def test_update_credential_answers_500_when_the_database_is_not_connected(): +def test_update_credential_answers_500_when_the_database_is_not_connected(credential_store): """The other rejection this handler raises must carry its own status too.""" - with patch("litellm.proxy.proxy_server.prisma_client", None): - response = _patch_credential( - "any-name", - {"credential_name": "any-name", "credential_values": {"api_key": "sk-x"}, "credential_info": {}}, - ) + credential_store(connected=False) + + response = _patch_credential( + "any-name", + {"credential_name": "any-name", "credential_values": {"api_key": "sk-x"}, "credential_info": {}}, + ) assert response.status_code == 500, f"rejected write answered {response.status_code}: {response.text}" -def test_update_credential_still_answers_200_on_a_successful_write(): +def test_update_credential_still_answers_200_on_a_successful_write(credential_store): """The fix must not turn a legitimate update into an error; the dashboard and the Playwright credentials spec both assert the success path.""" stored = CredentialItem( @@ -76,16 +105,96 @@ def test_update_credential_still_answers_200_on_a_successful_write(): credential_values={"api_key": "sk-old"}, credential_info={"custom_llm_provider": "openai"}, ) - with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( - "litellm.proxy.proxy_server.master_key", "sk-test-master" - ), patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository") as repository: - repository.return_value.find_by_name = AsyncMock(return_value=stored) - repository.return_value.update_by_name = AsyncMock(return_value=None) + credential_store(find_by_name=AsyncMock(return_value=stored), update_by_name=AsyncMock(return_value=None)) - response = _patch_credential( - "existing", - {"credential_name": "existing", "credential_values": {"api_key": "sk-new"}, "credential_info": {}}, - ) + response = _patch_credential( + "existing", + {"credential_name": "existing", "credential_values": {"api_key": "sk-new"}, "credential_info": {}}, + ) assert response.status_code == 200, response.text assert response.json()["success"] is True + + +def test_delete_credential_answers_404_when_the_credential_does_not_exist(credential_store): + """Regression: prisma's ``delete`` hands back None when the ``where`` clause matched no row + instead of raising, and the handler never looked. Deleting a name that was never stored + answered 200 "Credential deleted successfully", so an operator scripting cleanup could not + tell a real deletion from a typo.""" + credential_store(delete_by_name=AsyncMock(return_value=None)) + + response = _delete_credential("definitely-not-there") + + assert response.status_code == 404, f"delete of a missing credential answered {response.status_code}: {response.text}" + assert "definitely-not-there" in response.text + + +def test_delete_credential_still_answers_200_and_drops_the_credential_from_memory(credential_store): + """The fix must not turn a real deletion into an error, and the deleted credential must + stop being served from the in-memory list the proxy routes on.""" + stored = CredentialItem( + credential_name="doomed", + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "openai"}, + ) + survivor = CredentialItem( + credential_name="keeper", + credential_values={"api_key": "sk-keep"}, + credential_info={}, + ) + credential_store(in_memory=(stored, survivor), delete_by_name=AsyncMock(return_value=MagicMock())) + + response = _delete_credential("doomed") + + assert response.status_code == 200, response.text + assert response.json()["success"] is True + assert [credential.credential_name for credential in litellm.credential_list] == ["keeper"] + + +def test_delete_credential_leaves_a_credential_that_only_exists_in_memory_in_place(credential_store): + """A credential declared in the config yaml is never written to the table, so the delete + matches no row. Reporting success would be the same lie: it comes straight back on the next + proxy boot. ``PATCH /credentials/{name}`` already answers 404 for that credential.""" + config_only = CredentialItem( + credential_name="from-config-yaml", + credential_values={"api_key": "sk-config"}, + credential_info={}, + ) + credential_store(in_memory=(config_only,), delete_by_name=AsyncMock(return_value=None)) + + response = _delete_credential("from-config-yaml") + + assert response.status_code == 404, response.text + assert [credential.credential_name for credential in litellm.credential_list] == ["from-config-yaml"] + + +def test_delete_credential_answers_500_when_the_database_is_not_connected(credential_store): + """The handler used to ``return handle_exception_on_proxy(e)``, which makes the exception the + response body and lets FastAPI answer 200. A DB-less proxy answered its own 500 as a success.""" + credential_store(connected=False) + + response = _delete_credential("any-name") + + assert response.status_code == 500, f"rejected delete answered {response.status_code}: {response.text}" + + +class _CredentialThatCannotBeMasked: + """Stands in for anything that fails while ``GET /credentials`` builds its response.""" + + credential_name = "unreadable" + credential_info: dict = {} + + @property + def credential_values(self): + raise RuntimeError("credential store unreadable") + + +def test_get_credentials_answers_an_error_status_when_the_listing_fails(credential_store): + """Same ``return`` instead of ``raise`` on the list route: a failed listing was serialized as + a 200 whose body happened to be an error, so a caller reading the status saw an empty success.""" + credential_store(in_memory=(_CredentialThatCannotBeMasked(),)) + + response = _list_credentials() + + assert response.status_code == 500, f"failed listing answered {response.status_code}: {response.text}" + assert response.json().get("success") is not True From e80e78d3ef9a65db5cd2382d4cb7edc7b2554773 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 08:38:25 -0700 Subject: [PATCH 14/26] feat(cli): enable Claude Code gateway model discovery by default in lite claude (#39445) * feat(cli): enable Claude Code gateway model discovery by default in lite claude Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(cli): build agent env declaratively and document discovery key for lite up Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cli): keep build_agent_env within LIT002 type-discipline budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/client/cli/README.md | 6 +++--- litellm/proxy/client/cli/commands/agents.py | 8 +++++++- litellm/proxy/client/cli/commands/claude_settings.py | 11 +++++++++-- tests/test_litellm/proxy/client/cli/test_agents.py | 11 +++++++++++ .../test_litellm/proxy/client/cli/test_up_commands.py | 7 +++++++ 5 files changed, 37 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 3ddce35b53d..47355f328dd 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -489,7 +489,7 @@ lite codex exec "summarize the repo" Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process. -The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). +The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, and older versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). Options (these belong to the wrapper, so put them before the agent's own flags): @@ -505,7 +505,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_ ### Route Every Claude Code Session Through the Proxy -`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` when that key is missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` and `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY` to `1` when those keys are missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. @@ -529,7 +529,7 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi lite --base-url https://your-proxy.example.com login --config-claude ``` -It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. +It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, `env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`. diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index c591cbabee1..baa21996c7e 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -16,6 +16,8 @@ ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN" ANTHROPIC_API_KEY_ENV: Final = "ANTHROPIC_API_KEY" ENABLE_TOOL_SEARCH_ENV: Final = "ENABLE_TOOL_SEARCH" ENABLE_TOOL_SEARCH_VALUE: Final = "true" +ENABLE_GATEWAY_MODEL_DISCOVERY_ENV: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" +ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1" OPENAI_BASE_URL_ENV: Final = "OPENAI_BASE_URL" OPENAI_API_KEY_ENV: Final = "OPENAI_API_KEY" @@ -67,7 +69,9 @@ def build_agent_env( Anthropic key cannot win over the bearer token we set. ENABLE_TOOL_SEARCH defaults to true because Claude Code turns tool search off when ANTHROPIC_BASE_URL is not a first-party Anthropic host; a value already in - the environment is left alone. + the environment is left alone. CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY + defaults to 1 so Claude Code (v2.1.129+) fills its /model picker from the + proxy's /v1/models; likewise left alone when already set. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") @@ -77,6 +81,8 @@ def build_agent_env( env.pop(ANTHROPIC_API_KEY_ENV, None) if ENABLE_TOOL_SEARCH_ENV not in env: env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE + if ENABLE_GATEWAY_MODEL_DISCOVERY_ENV not in env: + env[ENABLE_GATEWAY_MODEL_DISCOVERY_ENV] = ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE if PROFILE_OPENAI in profiles: env[OPENAI_BASE_URL_ENV] = root + "/v1" env[OPENAI_API_KEY_ENV] = api_key diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 46af641636e..5e3ce95f088 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -26,6 +26,8 @@ ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" ENABLE_TOOL_SEARCH_VALUE: Final = "true" +ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: Final = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" +ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE: Final = "1" CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" @@ -77,13 +79,16 @@ def merge_claude_settings( stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH defaults to true because Claude Code turns tool search off when - ANTHROPIC_BASE_URL is not a first-party Anthropic host; an existing value is - left alone. Every other key is preserved untouched. + ANTHROPIC_BASE_URL is not a first-party Anthropic host, and + CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY defaults to 1 so the /model picker + is filled from the proxy's /v1/models; existing values of both are left + alone. Every other key is preserved untouched. """ raw_env: Final = settings.get(ENV_KEY, {}) base_env: Final = raw_env if isinstance(raw_env, dict) else {} env: Final = { ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, + ENABLE_GATEWAY_MODEL_DISCOVERY_KEY: ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE, **{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY}, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), } @@ -156,6 +161,8 @@ __all__ = ( "AUTOROUTE_BACKUP_PATH", "BACKUP_PATH", "CLAUDE_SETTINGS_PATH", + "ENABLE_GATEWAY_MODEL_DISCOVERY_KEY", + "ENABLE_GATEWAY_MODEL_DISCOVERY_VALUE", "ENABLE_TOOL_SEARCH_KEY", "ENABLE_TOOL_SEARCH_VALUE", "ENV_KEY", diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index 0191dad3d94..62b94e948be 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -78,9 +78,19 @@ class TestBuildAgentEnv: assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" assert env["ENABLE_TOOL_SEARCH"] == "true" + assert env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" assert "OPENAI_BASE_URL" not in env assert "OPENAI_API_KEY" not in env + def test_anthropic_profile_preserves_existing_gateway_model_discovery(self): + env = build_agent_env( + {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "0"}, + "http://localhost:4000", + "sk-key", + frozenset({"anthropic"}), + ) + assert env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "0" + def test_anthropic_profile_preserves_existing_tool_search(self): env = build_agent_env( {"ENABLE_TOOL_SEARCH": "false"}, @@ -107,6 +117,7 @@ class TestBuildAgentEnv: assert env["OPENAI_API_KEY"] == "sk-key" assert "ANTHROPIC_BASE_URL" not in env assert "ENABLE_TOOL_SEARCH" not in env + assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in env def test_both_profiles_set_everything(self): env = build_agent_env( diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index c78bdfa75b1..aead1764b0e 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -56,8 +56,14 @@ class TestMergeClaudeSettings: merged = merge_claude_settings(settings, "http://localhost:4000/", "new-helper") assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1" assert merged["apiKeyHelper"] == "new-helper" + def test_preserves_existing_gateway_model_discovery(self): + settings = {"env": {"CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "0"}} + merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + assert merged["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "0" + def test_preserves_existing_tool_search(self): settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} merged = merge_claude_settings(settings, "http://localhost:4000", "helper") @@ -73,6 +79,7 @@ class TestMergeClaudeSettings: assert merged["env"] == { "ANTHROPIC_BASE_URL": "http://localhost:4000", "ENABLE_TOOL_SEARCH": "true", + "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "1", } assert merged["apiKeyHelper"] == "helper" From 1de960bce7a4ae0ddd4bba419efc25044511f749 Mon Sep 17 00:00:00 2001 From: Rakesh <51485939+rakeshrepository@users.noreply.github.com> Date: Thu, 3 Sep 2026 21:16:07 +0530 Subject: [PATCH 15/26] fix(docker): bump nginx runtime to 1.31.5-alpine3.24 and pin digest to resolve critical CVEs (#39561) --- ui/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/Dockerfile b/ui/Dockerfile index 24140093270..f14b631685a 100644 --- a/ui/Dockerfile +++ b/ui/Dockerfile @@ -3,7 +3,7 @@ # UI container — Next.js static export served by nginx. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 -ARG NGINX_VERSION=1.31-alpine +ARG NGINX_VERSION=1.31.5-alpine3.24@sha256:34f40471dea485273c5e2a04dd5e97a682332ceb4a9adecd67de450dcb2fb390 # ---------- builder ---------- FROM ${UI_BUILD_IMAGE} AS builder From 34d4f7f8aef2951ffcf5ee04f33bff69aab4fd9d Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 08:58:48 -0700 Subject: [PATCH 16/26] fix: 1.99.0-rc2 UI bug batch (empty org on key create, session pagination, access group rename/delete) (#39436) * fix(ui): clearing the organization picker no longer sends organization_id="" on key create * fix(proxy): paginate Request Logs by conversation and aggregate session type counts and models server-side * fix(proxy): keep access groups in sync when a model is renamed or deleted * fix(proxy): cap the Request Logs conversation total like the row total * fix(proxy): judge access group backing by the database for db models A worker whose router has not polled the database yet still lists a sibling under its old name, so a delete or rename handled there kept the stale name in every access group. Only config-sourced deployments count as router backing now; db models are counted in the table. * fix(ui): keep the conversation badge when an MCP call represents a conversation A conversation that straddles the bounded page window can be represented by one of its MCP rows, which showed a plain MCP badge and hid the session counts. The badge now reads the server aggregates whenever the conversation has more than one call. * fix(proxy): list every model of a conversation in Request Logs and keep the conversation badge for MCP representatives Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): type session spend aggregates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): satisfy request logs lint budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): cap per-session model aggregation in request logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet type-discipline budget after staging merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): send an explicit null when the key edit form clears the organization Clearing the Organization picker in the key edit form wrote undefined into the form value, and JSON.stringify drops undefined-valued keys, so /key/update never saw the field and the key kept its old organization. Writing null instead survives serialization, and the backend's model_dump(exclude_unset=True) preserves it, so the column is set to NULL. --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 11 +- .../model_management_endpoints.py | 54 ++++- .../access_group_model_sync.py | 119 +++++++++++ .../spend_management_endpoints.py | 116 +++++++---- .../test_key_management_endpoints.py | 10 + .../test_model_management_endpoints.py | 188 ++++++++++++++++++ .../test_access_group_model_sync.py | 170 ++++++++++++++++ .../test_spend_management_endpoints.py | 51 +++++ type-discipline-budget.json | 2 +- .../create_key_button.integration.test.tsx | 13 ++ .../organisms/create_key_button.tsx | 2 +- .../templates/key_edit_view.test.tsx | 26 +++ .../components/templates/key_edit_view.tsx | 4 +- .../RequestLogsTableColumns.test.tsx | 62 ++++++ .../view_logs/RequestLogsTableColumns.tsx | 24 ++- .../src/components/view_logs/columns.tsx | 2 + 16 files changed, 797 insertions(+), 57 deletions(-) create mode 100644 litellm/proxy/management_helpers/access_group_model_sync.py create mode 100644 tests/test_litellm/proxy/management_helpers/test_access_group_model_sync.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 849e54c65aa..c6ecb0be8b9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1216,9 +1216,9 @@ class GenerateKeyRequest(KeyRequestBase): organization_id: str | None = None project_id: str | None = None - @field_validator("team_id", mode="before") + @field_validator("team_id", "organization_id", mode="before") @classmethod - def treat_cleared_team_id_as_unset(cls, v: object) -> object: + def treat_cleared_id_as_unset(cls, v: object) -> object: if v == "": return None return v @@ -1278,6 +1278,13 @@ class UpdateKeyRequest(KeyRequestBase): rotation_interval: str | None = None organization_id: str | None = None + @field_validator("organization_id", mode="before") + @classmethod + def treat_cleared_organization_id_as_unset(cls, v: object) -> object: + if v == "": + return None + return v + @model_validator(mode="after") def validate_temp_budget(self) -> "UpdateKeyRequest": if self.temp_budget_increase is not None or self.temp_budget_expiry is not None: diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ca66640bf46..613e726f89d 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -68,6 +68,10 @@ from litellm.proxy.management_endpoints.team_endpoints import ( from litellm.proxy.management_endpoints.team_endpoints import ( update_team as _legacy_update_team, ) +from litellm.proxy.management_helpers.access_group_model_sync import ( + sync_access_groups_for_deleted_model, + sync_access_groups_for_renamed_model, +) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log from litellm.proxy.spend_tracking.ptu_feature_flag import ( PTU_COST_ATTRIBUTION_ENV_VAR, @@ -715,6 +719,7 @@ async def patch_model( existing_params=db_model.litellm_params, ) + requested_model_name: Final = patch_data.model_name # Handle team model updates with proper alias management update_data: Final = await _update_team_model_in_db( db_model=db_model, @@ -741,6 +746,20 @@ async def patch_model( param=None, ) + stored_model_name: Final = update_data.get("model_name") + if ( + stored_model_name is not None + and stored_model_name == requested_model_name + and stored_model_name != db_model.model_name + ): + await sync_access_groups_for_renamed_model( + prisma_client=prisma_client, + model_id=model_id, + old_name=db_model.model_name, + new_name=stored_model_name, + llm_router=llm_router, + ) + # Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates) live_before_reload: Final = live_model_ids_snapshot() reload_outcome: Final = await clear_cache() @@ -1673,6 +1692,12 @@ async def delete_model( proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, ) + await sync_access_groups_for_deleted_model( + prisma_client=prisma_client, + model_id=model_info.id, + model_name=model_params.model_name, + llm_router=llm_router, + ) ## CREATE AUDIT LOG ## asyncio.create_task( @@ -2027,25 +2052,36 @@ async def update_model( model_params.litellm_params[k] = encrypted_value ### MERGE WITH EXISTING DATA ### - merged_dictionary: Final = {} _mp: Final[dict[str, object]] = model_params.litellm_params.dict() + merged_dictionary: Final = { + key: _existing_litellm_params_dict[key] if value is None else value + for key, value in _mp.items() + if value is not None or _existing_litellm_params_dict.get(key) is not None + } - for key, value in _mp.items(): - if value is not None: - merged_dictionary[key] = value - elif key in _existing_litellm_params_dict and _existing_litellm_params_dict[key] is not None: - merged_dictionary[key] = _existing_litellm_params_dict[key] - else: - pass - + renamed_to: Final = ( + model_params.model_name + if model_params.model_name not in (None, deployment.model_name) + and deployment.model_info.team_id is None + else None + ) _data: Final[dict[str, str]] = { "litellm_params": json.dumps(merged_dictionary), "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + **({} if renamed_to is None else {"model_name": renamed_to}), } model_response: Final = await _proxy_model_table(prisma_client).update( where={"model_id": _model_id}, data=_data, ) + if renamed_to is not None: + await sync_access_groups_for_renamed_model( + prisma_client=prisma_client, + model_id=_model_id, + old_name=deployment.model_name, + new_name=renamed_to, + llm_router=llm_router, + ) # Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates) live_before_reload: Final = live_model_ids_snapshot() diff --git a/litellm/proxy/management_helpers/access_group_model_sync.py b/litellm/proxy/management_helpers/access_group_model_sync.py new file mode 100644 index 00000000000..b9d81f2981f --- /dev/null +++ b/litellm/proxy/management_helpers/access_group_model_sync.py @@ -0,0 +1,119 @@ +""" +Keep `litellm_accessgrouptable.access_model_names` pointing at deployment names that still exist. + +Unified access groups store model names, not ids, so a deployment rename or delete that leaves +the arrays alone strands every group on a name nothing serves any more. +""" + +from collections.abc import Sequence +from typing import Final, Protocol + +from pydantic import BaseModel + +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient +from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_caches +from litellm.repositories.table_repositories import AccessGroupRepository +from litellm.router import Router + + +class _TouchedGroupRow(BaseModel): + access_group_id: str + + +class _DeploymentCountRow(BaseModel): + deployment_count: int + + +class _RawExecutor(Protocol): + async def query_raw(self, query: str, *args: str) -> Sequence[object]: ... + + +_BACKING_DEPLOYMENTS_SQL: Final = ( + 'SELECT COUNT(*)::int AS deployment_count FROM "LiteLLM_ProxyModelTable" WHERE "model_name" = $1' +) + +_REPLACE_MODEL_NAME_SQL: Final = ( + 'UPDATE "LiteLLM_AccessGroupTable" ' + 'SET "access_model_names" = array_replace(array_remove("access_model_names", $2), $1, $2) ' + 'WHERE $1 = ANY("access_model_names") ' + 'RETURNING "access_group_id"' +) + +_APPEND_MODEL_NAME_SQL: Final = ( + 'UPDATE "LiteLLM_AccessGroupTable" ' + 'SET "access_model_names" = array_append("access_model_names", $2) ' + 'WHERE $1 = ANY("access_model_names") AND NOT ($2 = ANY("access_model_names")) ' + 'RETURNING "access_group_id"' +) + +_REMOVE_MODEL_NAME_SQL: Final = ( + 'UPDATE "LiteLLM_AccessGroupTable" ' + 'SET "access_model_names" = array_remove("access_model_names", $1) ' + 'WHERE $1 = ANY("access_model_names") ' + 'RETURNING "access_group_id"' +) + + +def _raw_executor(prisma_client: object) -> _RawExecutor: + db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + return WriterPinnedClient(db).db # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin + + +def _config_sourced_sibling(llm_router: Router, deployment_id: str, model_id: str) -> bool: + if deployment_id == model_id: + return False + deployment: Final = llm_router.get_deployment(model_id=deployment_id) + return deployment is not None and not deployment.model_info.db_model + + +def _served_by_a_config_deployment(llm_router: Router | None, model_name: str, model_id: str) -> bool: + if llm_router is None: + return False + return any( + _config_sourced_sibling(llm_router, deployment_id, model_id) + for deployment_id in llm_router.get_model_ids(model_name=model_name) + ) + + +async def _still_backed(executor: _RawExecutor, llm_router: Router | None, model_name: str, model_id: str) -> bool: + if _served_by_a_config_deployment(llm_router, model_name, model_id): + return True + count_rows: Final = await executor.query_raw(_BACKING_DEPLOYMENTS_SQL, model_name) + return any(_DeploymentCountRow.model_validate(row).deployment_count > 0 for row in count_rows) + + +async def _rewrite_groups(executor: _RawExecutor, sql: str, *names: str) -> None: + touched_rows: Final = await executor.query_raw(sql, *names) + await invalidate_access_group_caches( + tuple(_TouchedGroupRow.model_validate(row).access_group_id for row in touched_rows) + ) + + +async def sync_access_groups_for_renamed_model( + prisma_client: object, + *, + model_id: str, + old_name: str, + new_name: str, + llm_router: Router | None, +) -> None: + if old_name == new_name: + return + executor: Final = _raw_executor(prisma_client) + old_name_still_backed: Final = await _still_backed(executor, llm_router, old_name, model_id) + await _rewrite_groups( + executor, _APPEND_MODEL_NAME_SQL if old_name_still_backed else _REPLACE_MODEL_NAME_SQL, old_name, new_name + ) + + +async def sync_access_groups_for_deleted_model( + prisma_client: object, + *, + model_id: str, + model_name: str, + llm_router: Router | None, +) -> None: + executor: Final = _raw_executor(prisma_client) + if await _still_backed(executor, llm_router, model_name, model_id): + return + await _rewrite_groups(executor, _REMOVE_MODEL_NAME_SQL, model_name) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index bb7dfafb297..dcc798b81cb 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -12,6 +12,7 @@ from typing import ( Literal, NamedTuple, Protocol, + TypeAlias, TypedDict, TypeVar, cast, # noqa: TID251 # prisma group_by returns untyped aggregate mappings @@ -19,6 +20,7 @@ from typing import ( import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response, status +from pydantic import TypeAdapter from typing_extensions import ReadOnly import litellm @@ -158,6 +160,26 @@ class _SessionSpendRow(TypedDict): session_cache_hit_count: ReadOnly[int] session_llm_count: ReadOnly[int] session_agent_count: ReadOnly[int] + session_models: ReadOnly[Sequence[str]] + + +_SESSION_MODELS_LIMIT: Final = 10 +_SESSION_MODEL_NAME_MAX_LEN: Final = 256 + + +class _SessionSpendStats(NamedTuple): + session_total_count: int + session_total_spend: float + mcp_tool_call_count: int + mcp_tool_call_spend: float + session_cache_hit_count: int + session_llm_count: int + session_agent_count: int + session_models: Sequence[str] + session_models_truncated: bool + + +_SessionSpendMap: TypeAlias = Mapping[tuple[str, str], _SessionSpendStats] class _SpendSumAggregate(TypedDict, total=False): @@ -4121,7 +4143,7 @@ async def _build_ui_spend_logs_response( } ) - session_spend_map: dict[tuple[str, str], dict[str, int | float]] = {} + session_spend_map: _SessionSpendMap = {} if enrich_session_counts and session_ids: from prisma.errors import PrismaError @@ -4139,40 +4161,60 @@ async def _build_ui_spend_logs_response( rows: Final[Sequence[_SessionSpendRow]] = await _query_raw( prisma_client, f""" - SELECT session_id, api_key, - COUNT(*)::int AS session_total_count, - COALESCE(SUM(spend), 0)::double precision AS session_total_spend, - COUNT(*) FILTER ( - WHERE call_type IN {_MCP_CALL_TYPES_SQL} - )::int AS mcp_tool_call_count, - COALESCE(SUM(spend) FILTER ( - WHERE call_type IN {_MCP_CALL_TYPES_SQL} - ), 0)::double precision AS mcp_tool_call_spend, - COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count, - COUNT(*) FILTER ( - WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL} - )::int AS session_llm_count, - COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count - FROM "LiteLLM_SpendLogs" - WHERE session_id = ANY($1::text[]) - AND api_key = ANY($2::text[]) - GROUP BY session_id, api_key + SELECT s.*, COALESCE(m.session_models, ARRAY[]::text[]) AS session_models + FROM ( + SELECT session_id, api_key, + COUNT(*)::int AS session_total_count, + COALESCE(SUM(spend), 0)::double precision AS session_total_spend, + COUNT(*) FILTER ( + WHERE call_type IN {_MCP_CALL_TYPES_SQL} + )::int AS mcp_tool_call_count, + COALESCE(SUM(spend) FILTER ( + WHERE call_type IN {_MCP_CALL_TYPES_SQL} + ), 0)::double precision AS mcp_tool_call_spend, + COUNT(*) FILTER (WHERE LOWER(cache_hit) = 'true')::int AS session_cache_hit_count, + COUNT(*) FILTER ( + WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL} + )::int AS session_llm_count, + COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count + FROM "LiteLLM_SpendLogs" + WHERE session_id = ANY($1::text[]) + AND api_key = ANY($2::text[]) + GROUP BY session_id, api_key + ) s + LEFT JOIN LATERAL ( + SELECT ARRAY_AGG(d.model ORDER BY d.model) AS session_models + FROM ( + SELECT DISTINCT LEFT(model, $3::int) AS model + FROM "LiteLLM_SpendLogs" + WHERE session_id = s.session_id + AND api_key = s.api_key + AND model IS NOT NULL AND model <> '' + ORDER BY 1 + LIMIT $4::int + ) d + ) m ON TRUE """, session_ids, authorized_api_keys, + _SESSION_MODEL_NAME_MAX_LEN, + _SESSION_MODELS_LIMIT + 1, ) session_spend_map = { - (row["session_id"], row["api_key"]): { - "session_total_count": int(row.get("session_total_count") or 0), - "session_total_spend": float(row.get("session_total_spend") or 0.0), - "mcp_tool_call_count": int(row.get("mcp_tool_call_count") or 0), - "mcp_tool_call_spend": float(row.get("mcp_tool_call_spend") or 0.0), - "session_cache_hit_count": int(row.get("session_cache_hit_count") or 0), - "session_llm_count": int(row.get("session_llm_count") or 0), - "session_agent_count": int(row.get("session_agent_count") or 0), - } + (row["session_id"], row["api_key"]): _SessionSpendStats( + session_total_count=int(row.get("session_total_count") or 0), + session_total_spend=float(row.get("session_total_spend") or 0.0), + mcp_tool_call_count=int(row.get("mcp_tool_call_count") or 0), + mcp_tool_call_spend=float(row.get("mcp_tool_call_spend") or 0.0), + session_cache_hit_count=int(row.get("session_cache_hit_count") or 0), + session_llm_count=int(row.get("session_llm_count") or 0), + session_agent_count=int(row.get("session_agent_count") or 0), + session_models=models[:_SESSION_MODELS_LIMIT], + session_models_truncated=len(models) > _SESSION_MODELS_LIMIT, + ) for row in rows if row.get("session_id") and row.get("api_key") is not None + for models in (TypeAdapter(list[str]).validate_python(row.get("session_models") or ()),) } except PrismaError: verbose_proxy_logger.debug( @@ -4187,15 +4229,17 @@ async def _build_ui_spend_logs_response( sid = row_dict.get("session_id") row_api_key = row_dict.get("api_key") session_stats = session_spend_map.get((sid, row_api_key)) if sid and row_api_key is not None else None - row_dict["session_total_count"] = int(session_stats["session_total_count"]) if session_stats else 1 + row_dict["session_total_count"] = session_stats.session_total_count if session_stats else 1 if session_stats: - row_dict["session_total_spend"] = session_stats["session_total_spend"] - if session_stats["mcp_tool_call_count"]: - row_dict["mcp_tool_call_count"] = session_stats["mcp_tool_call_count"] - row_dict["mcp_tool_call_spend"] = session_stats["mcp_tool_call_spend"] - row_dict["session_cache_hit_count"] = session_stats["session_cache_hit_count"] - row_dict["session_llm_count"] = session_stats["session_llm_count"] - row_dict["session_agent_count"] = session_stats["session_agent_count"] + row_dict["session_total_spend"] = session_stats.session_total_spend + if session_stats.mcp_tool_call_count: + row_dict["mcp_tool_call_count"] = session_stats.mcp_tool_call_count + row_dict["mcp_tool_call_spend"] = session_stats.mcp_tool_call_spend + row_dict["session_cache_hit_count"] = session_stats.session_cache_hit_count + row_dict["session_llm_count"] = session_stats.session_llm_count + row_dict["session_agent_count"] = session_stats.session_agent_count + row_dict["session_models"] = session_stats.session_models + row_dict["session_models_truncated"] = session_stats.session_models_truncated enriched.append(row_dict) response_data: list = enriched else: 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 7e2e680743f..68704f476b5 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 @@ -17505,6 +17505,16 @@ def test_generate_key_request_blank_team_id_is_personal(): assert GenerateKeyRequest(team_id="team-1").team_id == "team-1" +def test_key_request_blank_organization_id_is_unset(): + from litellm.proxy._types import RegenerateKeyRequest, UpdateKeyRequest + + assert GenerateKeyRequest(organization_id="").organization_id is None + assert RegenerateKeyRequest(organization_id="").organization_id is None + assert UpdateKeyRequest(key="sk-1", organization_id="").organization_id is None + assert GenerateKeyRequest(organization_id="org-1").organization_id == "org-1" + assert UpdateKeyRequest(key="sk-1", organization_id="org-1").organization_id == "org-1" + + def test_key_generation_check_blank_team_id_uses_personal_permissions(monkeypatch): """key_generation_check with team_id="" must take the personal-key path instead of failing the team lookup with "Unable to find team object" (LIT-3925).""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 4661cc17dbc..5fa59a85c9d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1,5 +1,6 @@ import inspect import asyncio +import contextlib import json from typing import Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -875,6 +876,7 @@ class TestDeleteModelClearsRouterRegistry: mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) @@ -936,6 +938,7 @@ class TestDeleteModelClearsRouterRegistry: mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) @@ -2079,6 +2082,7 @@ class TestAddAndDeleteModelLifecycle: mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( return_value=db_row @@ -2191,6 +2195,7 @@ class TestDeleteTeamBYOKModelGhost: mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( return_value=db_row ) @@ -2273,6 +2278,7 @@ class TestDeleteTeamBYOKModelGhost: mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( return_value=db_row ) @@ -2349,6 +2355,7 @@ class TestDeleteTeamBYOKModelGhost: mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( return_value=deleted_row ) @@ -2434,6 +2441,7 @@ class TestDeleteTeamBYOKModelGhost: mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( return_value=db_row ) @@ -2515,6 +2523,7 @@ class TestDeleteTeamBYOKModelGhost: mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( return_value=db_row ) @@ -2584,6 +2593,7 @@ class TestDeleteModelTeamAuth: mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( return_value=db_row ) @@ -2702,6 +2712,7 @@ class TestDeleteModelTeamAuth: mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.query_raw = AsyncMock(return_value=[]) mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( return_value=db_row ) @@ -3844,6 +3855,7 @@ class TestDeleteEvictionsHoldTheReconcileLock: prisma = MagicMock() prisma.db.litellm_proxymodeltable = table + prisma.db.query_raw = AsyncMock(return_value=[]) router = MagicMock() router.delete_deployment = MagicMock(return_value=True) @@ -4654,3 +4666,179 @@ class TestBlockModelResponseSerialization: assert body["model_id"] == "m-block-1" assert body["blocked"] is blocked assert body["litellm_params"] == {"model": "openai/gpt-4o-mini", "api_key": "encrypted-value"} + + +class TestAccessGroupModelSync: + """A rename or delete of a deployment must land in every unified access group that names it.""" + + _PS = "litellm.proxy.proxy_server" + _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" + _INVALIDATE = "litellm.proxy.management_helpers.access_group_model_sync.invalidate_access_group_caches" + + @staticmethod + def _admin(): + return UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin") + + @staticmethod + def _prisma_with_row(model_id: str, model_name: str, deployment_count: int): + row = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name=model_name, + litellm_params={"model": "openai/gpt-5.6"}, + model_info={"id": model_id}, + created_by="admin", + updated_by="admin", + ) + + async def query_raw(sql, *params): + if sql.startswith("SELECT COUNT(*)"): + return [{"deployment_count": deployment_count}] + return [{"access_group_id": "ag-1"}] + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=query_raw) + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=row) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=row) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=row) + return mock_prisma + + @staticmethod + def _access_group_updates(mock_prisma): + return [ + call + for call in mock_prisma.db.query_raw.await_args_list + if call.args[0].startswith('UPDATE "LiteLLM_AccessGroupTable"') + ] + + @contextlib.contextmanager + def _endpoint_env(self, mock_prisma, router): + with contextlib.ExitStack() as stack: + for target in ( + patch(f"{self._PS}.prisma_client", mock_prisma), + patch(f"{self._PS}.llm_router", router), + patch(f"{self._PS}.store_model_in_db", True), + patch(f"{self._PS}.premium_user", True), + patch(f"{self._PS}.proxy_logging_obj", MagicMock()), + patch(f"{self._PS}.user_api_key_cache", MagicMock()), + patch(f"{self._MOD}.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None)), + patch( + f"{self._MOD}.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + patch(f"{self._MOD}.encrypt_value_helper", side_effect=lambda value, **kwargs: value), + ): + stack.enter_context(target) + yield stack.enter_context(patch(self._INVALIDATE, new=AsyncMock())) + + @pytest.mark.asyncio + async def test_patch_model_rename_rewrites_the_groups_that_named_the_model(self): + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + mock_prisma = self._prisma_with_row("m-rename", "gpt-5.6", deployment_count=0) + router = MagicMock() + router.get_model_ids.return_value = ["m-rename"] + + with self._endpoint_env(mock_prisma, router) as invalidate: + await patch_model( + model_id="m-rename", + patch_data=updateDeployment(model_name="gpt-5.6-eu"), + user_api_key_dict=self._admin(), + ) + + written = mock_prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"] + assert written["model_name"] == "gpt-5.6-eu" + (update_call,) = self._access_group_updates(mock_prisma) + assert "array_replace" in update_call.args[0] + assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") + invalidate.assert_awaited_once_with(("ag-1",)) + + @pytest.mark.asyncio + async def test_patch_model_rename_appends_when_a_sibling_deployment_keeps_the_old_name(self): + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + mock_prisma = self._prisma_with_row("m-rename", "gpt-5.6", deployment_count=1) + router = MagicMock() + router.get_model_ids.return_value = ["m-rename"] + + with self._endpoint_env(mock_prisma, router): + await patch_model( + model_id="m-rename", + patch_data=updateDeployment(model_name="gpt-5.6-eu"), + user_api_key_dict=self._admin(), + ) + + (update_call,) = self._access_group_updates(mock_prisma) + assert "array_append" in update_call.args[0] + assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") + + @pytest.mark.asyncio + async def test_patch_model_without_a_rename_leaves_access_groups_alone(self): + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + mock_prisma = self._prisma_with_row("m-same", "gpt-5.6", deployment_count=0) + router = MagicMock() + router.get_model_ids.return_value = ["m-same"] + + with self._endpoint_env(mock_prisma, router) as invalidate: + await patch_model(model_id="m-same", patch_data=updateDeployment(blocked=True), user_api_key_dict=self._admin()) + + mock_prisma.db.query_raw.assert_not_awaited() + invalidate.assert_not_awaited() + + @pytest.mark.asyncio + async def test_delete_model_drops_the_name_from_groups_when_nothing_backs_it(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete, delete_model + + mock_prisma = self._prisma_with_row("m-doomed", "gpt-5.6", deployment_count=0) + router = MagicMock() + router.get_model_ids.return_value = [] + + with self._endpoint_env(mock_prisma, router) as invalidate: + await delete_model(model_info=ModelInfoDelete(id="m-doomed"), user_api_key_dict=self._admin()) + + (update_call,) = self._access_group_updates(mock_prisma) + assert "array_remove" in update_call.args[0] + assert update_call.args[1:] == ("gpt-5.6",) + invalidate.assert_awaited_once_with(("ag-1",)) + + @pytest.mark.asyncio + async def test_delete_model_keeps_the_name_while_a_sibling_deployment_backs_it(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete, delete_model + + mock_prisma = self._prisma_with_row("m-doomed", "gpt-5.6", deployment_count=1) + router = MagicMock() + router.get_model_ids.return_value = [] + + with self._endpoint_env(mock_prisma, router) as invalidate: + await delete_model(model_info=ModelInfoDelete(id="m-doomed"), user_api_key_dict=self._admin()) + + assert self._access_group_updates(mock_prisma) == [] + invalidate.assert_not_awaited() + + @pytest.mark.asyncio + async def test_update_model_persists_a_new_model_name_and_rewrites_the_groups(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_model + from litellm.types.router import ModelInfo, updateLiteLLMParams + + mock_prisma = self._prisma_with_row("m-terraform", "gpt-5.6", deployment_count=0) + router = MagicMock() + router.get_model_ids.return_value = ["m-terraform"] + + with self._endpoint_env(mock_prisma, router) as invalidate: + await update_model( + model_params=updateDeployment( + model_name="gpt-5.6-eu", + litellm_params=updateLiteLLMParams(model="openai/gpt-5.6"), + model_info=ModelInfo(id="m-terraform"), + ), + user_api_key_dict=self._admin(), + ) + + written = mock_prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"] + assert written["model_name"] == "gpt-5.6-eu" + (update_call,) = self._access_group_updates(mock_prisma) + assert "array_replace" in update_call.args[0] + assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") + invalidate.assert_awaited_once_with(("ag-1",)) diff --git a/tests/test_litellm/proxy/management_helpers/test_access_group_model_sync.py b/tests/test_litellm/proxy/management_helpers/test_access_group_model_sync.py new file mode 100644 index 00000000000..65ef2d55cb8 --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_access_group_model_sync.py @@ -0,0 +1,170 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy.db.prisma_client import PrismaWrapper +from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper +from litellm.proxy.management_helpers.access_group_model_sync import ( + sync_access_groups_for_deleted_model, + sync_access_groups_for_renamed_model, +) + +_INVALIDATE = "litellm.proxy.management_helpers.access_group_model_sync.invalidate_access_group_caches" + + +def _routed_prisma_client(deployment_count: int): + async def query_raw(sql, *params): + if sql.startswith("SELECT COUNT(*)"): + return [{"deployment_count": deployment_count}] + return [{"access_group_id": "ag-1"}, {"access_group_id": "ag-2"}] + + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + writer_inner.query_raw = AsyncMock(side_effect=query_raw) + reader_inner.query_raw = AsyncMock(side_effect=query_raw) + writer = PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False) + reader = PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + return SimpleNamespace(db=routing), writer_inner, reader_inner + + +def _access_group_updates(writer_inner): + return [ + call + for call in writer_inner.query_raw.await_args_list + if call.args[0].startswith('UPDATE "LiteLLM_AccessGroupTable"') + ] + + +@pytest.mark.asyncio +async def test_rename_replaces_the_old_name_when_no_other_deployment_carries_it(): + prisma_client, writer_inner, reader_inner = _routed_prisma_client(deployment_count=0) + + with patch(_INVALIDATE, new=AsyncMock()) as invalidate: + await sync_access_groups_for_renamed_model( + prisma_client, model_id="m-1", old_name="gpt-5.6", new_name="gpt-5.6-eu", llm_router=None + ) + + (update_call,) = _access_group_updates(writer_inner) + assert "array_replace" in update_call.args[0] + assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") + invalidate.assert_awaited_once_with(("ag-1", "ag-2")) + reader_inner.query_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_rename_appends_the_new_name_when_a_sibling_row_keeps_the_old_one(): + prisma_client, writer_inner, _ = _routed_prisma_client(deployment_count=1) + + with patch(_INVALIDATE, new=AsyncMock()) as invalidate: + await sync_access_groups_for_renamed_model( + prisma_client, model_id="m-1", old_name="gpt-5.6", new_name="gpt-5.6-eu", llm_router=None + ) + + (update_call,) = _access_group_updates(writer_inner) + assert "array_append" in update_call.args[0] + assert "array_replace" not in update_call.args[0] + assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") + invalidate.assert_awaited_once_with(("ag-1", "ag-2")) + + +def _router_serving(db_model_by_deployment_id: dict[str, bool]): + llm_router = MagicMock() + llm_router.get_model_ids.return_value = list(db_model_by_deployment_id) + llm_router.get_deployment.side_effect = lambda model_id: SimpleNamespace( + model_info=SimpleNamespace(db_model=db_model_by_deployment_id[model_id]) + ) + return llm_router + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_model_by_deployment_id, expected_write", + [ + ({"m-1": True}, "array_replace"), + ({"m-1": True, "m-from-config": False}, "array_append"), + ({"m-1": True, "m-db-sibling-this-worker-has-not-refreshed": True}, "array_replace"), + ], +) +async def test_rename_counts_only_config_deployments_with_another_id_as_backing_the_old_name( + db_model_by_deployment_id, expected_write +): + prisma_client, writer_inner, _ = _routed_prisma_client(deployment_count=0) + llm_router = _router_serving(db_model_by_deployment_id) + + with patch(_INVALIDATE, new=AsyncMock()): + await sync_access_groups_for_renamed_model( + prisma_client, model_id="m-1", old_name="gpt-5.6", new_name="gpt-5.6-eu", llm_router=llm_router + ) + + llm_router.get_model_ids.assert_called_once_with(model_name="gpt-5.6") + (update_call,) = _access_group_updates(writer_inner) + assert expected_write in update_call.args[0] + + +@pytest.mark.asyncio +async def test_delete_ignores_a_db_sibling_this_worker_has_not_refreshed_yet(): + prisma_client, writer_inner, _ = _routed_prisma_client(deployment_count=0) + llm_router = _router_serving({"m-1": True, "m-renamed-elsewhere": True}) + + with patch(_INVALIDATE, new=AsyncMock()) as invalidate: + await sync_access_groups_for_deleted_model( + prisma_client, model_id="m-1", model_name="gpt-5.6", llm_router=llm_router + ) + + (update_call,) = _access_group_updates(writer_inner) + assert "array_remove" in update_call.args[0] + invalidate.assert_awaited_once_with(("ag-1", "ag-2")) + + +@pytest.mark.asyncio +async def test_delete_keeps_the_name_while_a_config_deployment_still_serves_it(): + prisma_client, writer_inner, _ = _routed_prisma_client(deployment_count=0) + llm_router = _router_serving({"m-1": True, "m-from-config": False}) + + with patch(_INVALIDATE, new=AsyncMock()) as invalidate: + await sync_access_groups_for_deleted_model( + prisma_client, model_id="m-1", model_name="gpt-5.6", llm_router=llm_router + ) + + assert _access_group_updates(writer_inner) == [] + invalidate.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_rename_to_the_same_name_writes_nothing(): + prisma_client, writer_inner, _ = _routed_prisma_client(deployment_count=0) + + with patch(_INVALIDATE, new=AsyncMock()) as invalidate: + await sync_access_groups_for_renamed_model( + prisma_client, model_id="m-1", old_name="gpt-5.6", new_name="gpt-5.6", llm_router=None + ) + + assert _access_group_updates(writer_inner) == [] + invalidate.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_removes_the_name_when_no_row_backs_it_any_more(): + prisma_client, writer_inner, reader_inner = _routed_prisma_client(deployment_count=0) + + with patch(_INVALIDATE, new=AsyncMock()) as invalidate: + await sync_access_groups_for_deleted_model(prisma_client, model_id="m-1", model_name="gpt-5.6", llm_router=None) + + (update_call,) = _access_group_updates(writer_inner) + assert "array_remove" in update_call.args[0] + assert update_call.args[1:] == ("gpt-5.6",) + invalidate.assert_awaited_once_with(("ag-1", "ag-2")) + reader_inner.query_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_keeps_the_name_while_a_sibling_row_still_backs_it(): + prisma_client, writer_inner, _ = _routed_prisma_client(deployment_count=2) + + with patch(_INVALIDATE, new=AsyncMock()) as invalidate: + await sync_access_groups_for_deleted_model(prisma_client, model_id="m-1", model_name="gpt-5.6", llm_router=None) + + assert _access_group_updates(writer_inner) == [] + invalidate.assert_not_awaited() diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 30b086bab61..c8b35e8a841 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3971,6 +3971,7 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): "mcp_tool_call_spend": 10.0, "session_llm_count": 1, "session_agent_count": 0, + "session_models": ["claude-haiku-4-5", "gpt-5.4-nano"], } ] ) @@ -3997,6 +3998,7 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): assert rows[1]["mcp_tool_call_spend"] == 10.0 assert rows[0]["session_llm_count"] == 1 assert rows[0]["session_agent_count"] == 0 + assert rows[0]["session_models"] == ["claude-haiku-4-5", "gpt-5.4-nano"] # Every row in the session carries the full session spend, not just its own assert rows[0]["session_total_spend"] == 15.0 @@ -4004,11 +4006,60 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): # Row without a session_id defaults to 1 assert rows[2]["session_total_count"] == 1 + assert "session_models" not in rows[2] # The count is folded into the single aggregate query; no separate group_by call. mock_prisma.db.litellm_spendlogs.group_by.assert_not_called() +@pytest.mark.asyncio +async def test_build_ui_spend_logs_response_caps_session_models(): + """The per-session model list is bounded server-side and flags when it was cut.""" + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + _SESSION_MODELS_LIMIT, + _build_ui_spend_logs_response, + ) + + session_id = "sess-many-models" + api_key = "hashed-key-xyz" + over_limit_models = [f"model-{i:02d}" for i in range(_SESSION_MODELS_LIMIT + 1)] + + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "session_id": session_id, + "api_key": api_key, + "session_total_count": len(over_limit_models), + "session_total_spend": 1.0, + "mcp_tool_call_count": 0, + "mcp_tool_call_spend": 0.0, + "session_llm_count": len(over_limit_models), + "session_agent_count": 0, + "session_models": over_limit_models, + } + ] + ) + + result = await _build_ui_spend_logs_response( + prisma_client=mock_prisma, + data=[{"request_id": "req-1", "session_id": session_id, "call_type": "completion", "api_key": api_key}], + total_records=1, + page=1, + page_size=50, + total_pages=1, + enrich_session_counts=True, + ) + + row = result["data"][0] + assert row["session_models"] == over_limit_models[:_SESSION_MODELS_LIMIT] + assert row["session_models_truncated"] is True + + sql, *params = mock_prisma.db.query_raw.await_args.args + assert "LIMIT $4" in sql + assert params[3] == _SESSION_MODELS_LIMIT + 1 + + @pytest.mark.asyncio async def test_build_ui_spend_logs_response_key_split_session_gets_per_key_aggregates(): """ diff --git a/type-discipline-budget.json b/type-discipline-budget.json index d5bf3883be4..cbcb5dca443 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22330 + "limit": 22328 }, "LIT002": { "limit": 26763 diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 8630be8548f..045dcfa3ceb 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -750,6 +750,19 @@ describe("CreateKey", () => { expect((await createdPayload()).organization_id).toBe("org-1"); }); + + it("drops organization_id when the chosen organization is cleared again", async () => { + state.organizations = [{ organization_id: "org-1", organization_alias: "Engineering" }]; + await openModal(); + await nameTheKey(); + + await userEvent.click(await screen.findByLabelText("Organization")); + await userEvent.click(await screen.findByRole("option", { name: /Engineering/ })); + await userEvent.click(await screen.findByRole("button", { name: "Clear" })); + await submit(); + + expect((await createdPayload()).organization_id).toBeUndefined(); + }); }); describe("policy and prompt fields", () => { diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index d38a8995c1b..5749541dcea 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -588,7 +588,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp }; const changeOrganization = (write: FieldWrite) => (orgId: string) => { - write(orgId); + write(orgId || undefined); setSelectedOrganizationId(orgId || null); // Clear team and project when org changes setSelectedCreateKeyTeam(null); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index 97b09e00808..10b1983b54b 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -1457,6 +1457,32 @@ describe("KeyEditView", () => { expect(screen.getByLabelText("Organization")).toHaveValue("Engineering"); }); }); + + it("submits organization_id as null after the organization is cleared", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderWithProviders( + {}} + onSubmit={onSubmit} + accessToken="" + userID="" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByLabelText("Organization")).toHaveValue("Engineering"); + }); + await userEvent.click(screen.getByRole("button", { name: "Clear" })); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalledWith(expect.objectContaining({ organization_id: null })); + }); + expect(JSON.parse(JSON.stringify(onSubmit.mock.calls[0][0]))).toHaveProperty("organization_id", null); + }); }); describe("models dropdown team gating", () => { diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index df1af2ca8e9..3e772fd0e9b 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -303,8 +303,8 @@ export function KeyEditView({ } }; - const handleOrganizationChange = (setField: (value: string | undefined) => void, orgId: string | undefined) => { - setField(orgId); + const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | undefined) => { + setField(orgId || null); setSelectedOrganizationId(orgId || null); form.setValue("team_id", undefined); }; diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index ce59c62f1c4..e3bacc0908a 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -75,6 +75,68 @@ describe("Cost column", () => { }); }); +describe("Type column", () => { + it("shows the conversation badge and composition even when an MCP call represents the conversation", async () => { + const user = userEvent.setup(); + const mcpRepresentative = { + request_id: "req-mcp-rep", + call_type: "call_mcp_tool", + session_id: "sess-edge", + session_total_count: 3, + session_llm_count: 2, + mcp_tool_call_count: 1, + session_agent_count: 0, + }; + renderRows([logEntry(mcpRepresentative)]); + + expect(screen.queryByText("MCP")).not.toBeInTheDocument(); + await user.hover(screen.getByText("3")); + expect(await screen.findByText("2 LLM • 1 MCP")).toBeInTheDocument(); + }); + + it("keeps the plain MCP badge for a single MCP call", () => { + renderRows([logEntry({ request_id: "req-mcp-solo", call_type: "call_mcp_tool", session_total_count: 1 })]); + + expect(screen.getByText("MCP")).toBeInTheDocument(); + }); +}); + +describe("Model column", () => { + it("lists every model used across a conversation, not only the representative call's model", () => { + const conversationCall: Partial = { + request_id: "req-session", + model: "gpt-5.6", + session_id: "sess-1", + session_total_count: 3, + session_models: ["claude-sonnet-5", "gpt-5.6"], + }; + renderRows([logEntry(conversationCall)]); + + expect(screen.getByText("claude-sonnet-5, gpt-5.6")).toBeInTheDocument(); + expect(screen.queryByText("gpt-5.6")).not.toBeInTheDocument(); + }); + + it("marks a conversation whose model list was capped by the server", () => { + const cappedCall = { + request_id: "req-capped", + model: "gpt-5.6", + session_id: "sess-2", + session_total_count: 30, + session_models: ["claude-sonnet-5", "gpt-5.6"], + session_models_truncated: true, + }; + renderRows([logEntry(cappedCall)]); + + expect(screen.getByText("claude-sonnet-5, gpt-5.6, ...")).toBeInTheDocument(); + }); + + it("keeps a single call's own model", () => { + renderRows([logEntry({ request_id: "req-single", model: "gpt-5.6" })]); + + expect(screen.getByText("gpt-5.6")).toBeInTheDocument(); + }); +}); + describe("row action cells", () => { it("reports the key hash through the injected dependency rather than a row field", async () => { const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index b9058d02a6b..8db0b106851 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -63,9 +63,11 @@ export const getRequestLogsTableColumns = ({ const sessionAgentCount = log.session_agent_count ?? (isAgent ? sessionCount : 0); const sessionMcpCount = log.mcp_tool_call_count ?? (isMcp ? sessionCount : 0); - if (isMcp) return ; - if (isAgent && sessionCount <= 1) return ; - if (sessionCount <= 1) return ; + if (sessionCount <= 1) { + if (isMcp) return ; + if (isAgent) return ; + return ; + } const sessionTypeBadge = ( @@ -224,10 +226,13 @@ export const getRequestLogsTableColumns = ({ cell: ({ row }) => { const log = row.original; const provider = log.custom_llm_provider; - const modelName = log.model ?? ""; + const sessionModels = log.session_models ?? []; + const modelNames = sessionModels.length > 0 ? sessionModels : [log.model ?? ""]; + const modelLabel = log.session_models_truncated ? `${modelNames.join(", ")}, ...` : modelNames.join(", "); + const isSingleModel = modelNames.length === 1; return (
- {provider && ( + {provider && isSingleModel && ( )} - {modelName}} /> + + {modelLabel} + + } + />
); }, diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx index 2f3a3681352..21e09faf454 100644 --- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx @@ -47,4 +47,6 @@ export type LogEntry = { mcp_tool_call_spend?: number; session_llm_count?: number; session_agent_count?: number; + session_models?: string[]; + session_models_truncated?: boolean; }; From 4990f06accc847a886e2dfc2236a89abd8bd3acc Mon Sep 17 00:00:00 2001 From: moe-berri Date: Thu, 3 Sep 2026 08:59:04 -0700 Subject: [PATCH 17/26] feat(auto-router): support classifier reasoning effort (#39372) * feat(auto-router): support classifier reasoning effort * fix(auto-router): harden classifier reasoning effort * fix(ui): satisfy classifier config lint limits * refactor(auto-router): simplify classifier effort support * fix(auto-router): clear frontend-lint and type-discipline gates, trim LOC --------- Co-authored-by: Tin Chi Lo --- litellm/router.py | 119 ++++++++++++++-- .../complexity_router/README.md | 7 +- .../complexity_router/complexity_router.py | 14 +- .../complexity_router/config.py | 8 ++ .../router_strategy/test_complexity_router.py | 51 ++++++- tests/test_litellm/test_router.py | 130 +++++++++++++++++- .../add_model/ClassificationMethodConfig.tsx | 35 ++++- .../ClassifierReasoningEffortSelect.tsx | 86 ++++++++++++ .../add_model/ComplexityRouterConfig.test.tsx | 112 ++++++++++++++- .../add_model/ComplexityRouterConfig.tsx | 17 +-- .../add_model/HeuristicScoringConfig.test.tsx | 6 +- .../add_model/add_auto_router_tab.tsx | 35 +++-- .../auto_router_connection_test.test.tsx | 8 +- .../add_model/auto_router_connection_test.tsx | 10 +- .../build_auto_router_test_targets.test.ts | 28 ++++ .../build_auto_router_test_targets.ts | 17 ++- .../build_complexity_router_config.test.ts | 51 ++++++- .../build_complexity_router_config.ts | 33 ++++- .../add_model/complexity_router_tiers.ts | 18 +++ ...d_updated_complexity_router_config.test.ts | 6 +- .../edit_auto_router_modal.tsx | 7 + .../llm_calls/fetch_models.test.tsx | 18 +++ .../src/components/llm_calls/fetch_models.tsx | 6 +- .../src/components/networking.test.ts | 9 ++ .../src/components/networking.tsx | 6 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 26 files changed, 774 insertions(+), 68 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/ClassifierReasoningEffortSelect.tsx diff --git a/litellm/router.py b/litellm/router.py index 303b22c9484..2b8b342d253 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -50,6 +50,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, + INTERNAL_CALL_ORIGIN_METADATA_KEY, RUNTIME_UPDATABLE_ROUTER_SETTINGS, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) @@ -231,6 +232,7 @@ from litellm.types.router import ( ) from litellm.types.services import ServiceTypes from litellm.types.utils import ( + AUTOROUTER_CLASSIFIER_CALL_ORIGIN, PROMPT_QUOTING_ROUTING_DECISION_FIELDS, CustomPricingLiteLLMParams, GenericBudgetConfigType, @@ -2168,6 +2170,76 @@ class Router: verbose_router_logger.debug("Error occurred while printing deployment - %s", e) raise e + @staticmethod + def _deployment_params_with_request_reasoning_override( + deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object] + ) -> dict[str, object]: # mutable-ok: litellm's request pipeline consumes a mutable kwargs mapping + """Return deployment params whose equivalent effort controls cannot outrank a request override. + + Providers expose the same setting through several native carriers. A request-level + ``reasoning_effort`` is the portable override, so a deployment's ``thinking`` or nested + ``*.effort`` must not remain beside it and either win or trigger a conflicting-params 400. + Every changed mapping is copied so the Router's shared deployment config stays immutable. + """ + sanitized: Final = dict(deployment_params) # mutable-ok: request-local copy protects shared Router state + if request_kwargs.get("reasoning_effort") is None: + return sanitized + + sanitized.pop("thinking", None) + Router._pop_effort_from_nested_carrier(sanitized, "output_config") + Router._pop_effort_from_nested_carrier(sanitized, "reasoning") + + extra_body: Final = sanitized.get("extra_body") + if isinstance(extra_body, Mapping): + sanitized_extra_body: Final = dict(extra_body) # mutable-ok: request-local nested copy + sanitized_extra_body.pop("reasoning_effort", None) + sanitized_extra_body.pop("thinking", None) + Router._pop_effort_from_nested_carrier(sanitized_extra_body, "output_config") + Router._pop_effort_from_nested_carrier(sanitized_extra_body, "reasoning") + if sanitized_extra_body: + sanitized["extra_body"] = sanitized_extra_body + else: + sanitized.pop("extra_body", None) + return sanitized + + @staticmethod + def _is_classifier_internal_call(kwargs: Mapping[str, object]) -> bool: + metadata: Final = kwargs.get("metadata") + litellm_metadata: Final = kwargs.get("litellm_metadata") + return any( + isinstance(candidate, Mapping) + and candidate.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == AUTOROUTER_CLASSIFIER_CALL_ORIGIN + for candidate in (metadata, litellm_metadata) + ) + + def _drop_unsupported_classifier_reasoning_effort( + self, + deployment: DeploymentTypedDict, + model: str, + kwargs: dict[str, object], # mutable-ok: fallback must update the active request and its log body together + ) -> None: + """Let a classifier fallback without reasoning support remain a usable fallback. + + The dashboard only offers explicitly advertised levels, but an existing config can outlive + a model change and fallbacks can target a different group. Unknown capability fails open; + only a provider that explicitly rejects the parameter has it removed. + """ + if kwargs.get("reasoning_effort") is None or not self._is_classifier_internal_call(kwargs): + return + if self._deployment_accepts_param(deployment, model, "reasoning_effort"): + return + verbose_router_logger.warning( + "litellm.router.py: dropping classifier reasoning_effort for model=%s because the selected deployment does not support it", + model, + ) + kwargs.pop("reasoning_effort", None) + proxy_server_request: Final = kwargs.get("proxy_server_request") + if not isinstance(proxy_server_request, dict): + return + body: Final = proxy_server_request.get("body") + if isinstance(body, dict): + body.pop("reasoning_effort", None) + ### COMPLETION, EMBEDDING, IMG GENERATION FUNCTIONS def completion(self, model: str, messages: list[dict[str, str]], **kwargs) -> ModelResponse | CustomStreamWrapper: @@ -2203,9 +2275,16 @@ class Router: specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, ) + self._drop_unsupported_classifier_reasoning_effort( + deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment + model=model, + kwargs=kwargs, + ) # Check for silent model experiment # Make a local copy of litellm_params to avoid mutating the Router's state - litellm_params: Final = deployment["litellm_params"].copy() + litellm_params: Final = self._deployment_params_with_request_reasoning_override( + deployment["litellm_params"], kwargs + ) silent_model: Final = litellm_params.pop("silent_model", None) if silent_model is not None: @@ -3216,6 +3295,11 @@ class Router: specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, ) + self._drop_unsupported_classifier_reasoning_effort( + deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment + model=model, + kwargs=kwargs, + ) _timeout_debug_deployment_dict = deployment end_time: Final = time.time() @@ -3237,7 +3321,9 @@ class Router: # Check for silent model experiment # Make a local copy of litellm_params to avoid mutating the Router's state - litellm_params: Final = deployment["litellm_params"].copy() + litellm_params: Final = self._deployment_params_with_request_reasoning_override( + deployment["litellm_params"], kwargs + ) silent_model: Final = litellm_params.pop("silent_model", None) if silent_model is not None: @@ -10255,6 +10341,8 @@ class Router: total_itpm: int | None = None total_otpm: int | None = None configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None + reasoning_efforts_initialized = False + reasoning_efforts_unknown = False model_list: Final = self.get_model_list(model_name=model_group) if model_list is None: return None @@ -10441,10 +10529,23 @@ class Router: if model_info.get("rpm", None) is not None and _deployment_rpm is None: _deployment_rpm = model_info.get("rpm") - model_group_info.supported_reasoning_efforts = intersect_supported_reasoning_efforts( - model_group_info.supported_reasoning_efforts, - resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=deployment_is_mapped), + deployment_reasoning_efforts = ( + resolve_supported_reasoning_efforts( # rebind-ok: recalculated per deployment + model_info, deployment_is_mapped=deployment_is_mapped + ) ) + if deployment_reasoning_efforts is None: + reasoning_efforts_unknown = True + model_group_info.supported_reasoning_efforts = None + elif not reasoning_efforts_initialized: + reasoning_efforts_initialized = True + if not reasoning_efforts_unknown: + model_group_info.supported_reasoning_efforts = deployment_reasoning_efforts + elif not reasoning_efforts_unknown: + model_group_info.supported_reasoning_efforts = intersect_supported_reasoning_efforts( + model_group_info.supported_reasoning_efforts, + deployment_reasoning_efforts, + ) if _deployment_tpm is not None: if total_tpm is None: @@ -12308,10 +12409,12 @@ class Router: @staticmethod def _pop_effort_from_nested_carrier(request_kwargs: dict[str, object], carrier: str) -> None: nested: Final = request_kwargs.get(carrier) - if not isinstance(nested, dict): + if not isinstance(nested, Mapping): return - nested.pop("effort", None) - if not nested: + sanitized: Final = {key: value for key, value in nested.items() if key != "effort"} + if sanitized: + request_kwargs[carrier] = sanitized # rebind-ok: copy-on-write, so a shared nested carrier is never edited + else: request_kwargs.pop(carrier, None) @staticmethod diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index ee51add1ca1..e3da70f50fe 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -255,7 +255,8 @@ model_list: classifier_type: heuristic_first heuristic_first_max_tier: SIMPLE classifier_llm_config: - model: gpt-4o-mini + model: gpt-5-mini + reasoning_effort: low tiers: SIMPLE: gpt-4o-mini MEDIUM: gpt-4o @@ -263,6 +264,10 @@ model_list: REASONING: o1-preview ``` +`classifier_llm_config.reasoning_effort` applies only to the internal classifier call. Omit it to +keep the classifier deployment or provider default, or set a supported value such as `none` or +`low` to override that call. + A request short-circuits, meaning it routes on the scorer's own tier with no classifier call, when two things hold: the scorer landed at or below `heuristic_first_max_tier`, and it produced at least one signal. Everything else goes to the classifier, which then decides as it normally would. diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a00ae6bee80..17e3d1256d0 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -28,6 +28,7 @@ from pydantic import BaseModel, create_model from litellm._logging import verbose_router_logger from litellm.constants import ( EMPTY_MAPPING, + INTERNAL_CALL_ORIGIN_METADATA_KEY, RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, ) @@ -42,6 +43,7 @@ from litellm.router_strategy.complexity_router.tier_predictor import ( TierSuccessPredictor, resolve_tier_artifact, ) +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, ModelResponse, @@ -1668,20 +1670,27 @@ class ComplexityRouter(CustomLogger): ) request_metadata = (request_kwargs or {}).get("litellm_metadata") or (request_kwargs or {}).get("metadata") - metadata: Final = forwarded_internal_call_metadata(request_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN) + metadata: Final = { # mutable-ok: SDK metadata kwarg is enriched by the request pipeline + **forwarded_internal_call_metadata(request_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN), + INTERNAL_CALL_ORIGIN_METADATA_KEY: AUTOROUTER_CLASSIFIER_CALL_ORIGIN, + } turn_off_message_logging: Final = _effective_turn_off_message_logging(request_kwargs) - messages_for_call: Final = [ + messages_for_call: Final[list[AllMessageValues]] = [ # mutable-ok: SDK request payload list is built once {"role": "system", "content": classifier_system_prompt}, {"role": "user", "content": user_payload}, ] response_format: Final = classifier_response_format + classifier_call_params: Mapping[str, str] = EMPTY_MAPPING + if llm_config.reasoning_effort is not None: + classifier_call_params = MappingProxyType({"reasoning_effort": llm_config.reasoning_effort}) proxy_server_request: Final = { "body": { "model": llm_config.model, "messages": messages_for_call, "response_format": response_format, + **classifier_call_params, } } @@ -1693,6 +1702,7 @@ class ComplexityRouter(CustomLogger): metadata=metadata, proxy_server_request=proxy_server_request, turn_off_message_logging=turn_off_message_logging, + **classifier_call_params, **_parent_session_kwargs(request_kwargs), ) content: Final = response.choices[0].message.content diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 9f2054dda01..70c1b281e31 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -12,6 +12,7 @@ from typing import Annotated, Final, Literal from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator +from litellm.types.llms.openai import REASONING_EFFORT from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin from .tier_predictor import TrainedTierArtifact @@ -432,6 +433,13 @@ class ClassifierLLMConfig(BaseModel): model: str = Field( description="Model name (from the router's model_list) to call for classification", ) + reasoning_effort: REASONING_EFFORT | None = Field( + default=None, + description=( + "Reasoning effort override for classifier calls. Leave unset to use " + "the classifier deployment or provider default." + ), + ) timeout_ms: int = Field( default=3000, description="Timeout budget for the classification call, in milliseconds", diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 3ecccb673f3..aa1b51afe10 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -1558,6 +1558,14 @@ class TestLLMClassifierConfig: assert config.classifier_type == "heuristic" assert config.classifier_llm_config is None + @pytest.mark.parametrize("reasoning_effort", ["", "ultra"]) + def test_classifier_reasoning_effort_rejects_unsupported_values(self, reasoning_effort): + with pytest.raises(ValidationError): + ComplexityRouterConfig( + classifier_type="llm", + classifier_llm_config={"model": "haiku-classifier", "reasoning_effort": reasoning_effort}, + ) + CUSTOM_TIER_LABELS: Dict[str, str] = { "SIMPLE": "Cheap", @@ -1871,6 +1879,19 @@ class TestLLMClassifier: call_kwargs = mock_router_instance.acompletion.call_args.kwargs assert call_kwargs["metadata"] == {**request_metadata, "internal_call_origin": "autorouter_classifier"} + @pytest.mark.asyncio + async def test_aclassify_stamps_internal_origin_without_caller_metadata( + self, llm_complexity_router, mock_router_instance + ): + """Fallback handling must still recognize the classifier when an SDK caller supplied no metadata.""" + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) + + await llm_complexity_router.aclassify("hi") + + assert mock_router_instance.acompletion.call_args.kwargs["metadata"] == { + "internal_call_origin": "autorouter_classifier" + } + @pytest.mark.asyncio @pytest.mark.parametrize( "request_kwargs", @@ -1948,6 +1969,33 @@ class TestLLMClassifier: "REASONING", ] + @pytest.mark.asyncio + @pytest.mark.parametrize("reasoning_effort", [None, "none", "low"], ids=["omitted", "none", "low"]) + async def test_classifier_reasoning_effort_reaches_only_classifier_call( + self, mock_router_instance, llm_classifier_config, reasoning_effort + ): + classifier_llm_config = { + **llm_classifier_config["classifier_llm_config"], + **({"reasoning_effort": reasoning_effort} if reasoning_effort is not None else {}), + } + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**llm_classifier_config, "classifier_llm_config": classifier_llm_config}, + ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')) + + await router.aclassify("explain quantum tunneling in depth") + + call_kwargs = mock_router_instance.acompletion.call_args.kwargs + body = call_kwargs["proxy_server_request"]["body"] + if reasoning_effort is None: + assert "reasoning_effort" not in call_kwargs + assert "reasoning_effort" not in body + else: + assert call_kwargs["reasoning_effort"] == reasoning_effort + assert body["reasoning_effort"] == reasoning_effort + @pytest.mark.asyncio async def test_aclassify_propagates_top_level_turn_off_message_logging( self, llm_complexity_router, mock_router_instance @@ -8735,9 +8783,10 @@ class TestClassificationRubrics: [ {"model": "haiku-classifier", "system_prompt": "Grade the data sensitivity of the request."}, {"model": "haiku-classifier", "classification_rubric": "chat"}, + {"model": "haiku-classifier", "reasoning_effort": "low"}, {"model": "haiku-classifier"}, ], - ids=["custom-prompt", "chat-preset", "neither"], + ids=["custom-prompt", "chat-preset", "reasoning-effort", "neither"], ) def test_config_survives_a_dump_and_rebuild(self, classifier_llm_config): """/auto_router/test_routing dumps this config and hands the dict straight back to diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index c843a66a1c1..3b4c80b6b4f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -35,6 +35,7 @@ from litellm.router import ( _anthropic_stream_should_drop_pre_content_ping, _is_retriable_anthropic_status, ) +from litellm.types.router import DeploymentTypedDict def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -9783,11 +9784,12 @@ def test_model_group_info_intersects_supported_reasoning_efforts(): assert result.supported_reasoning_efforts == ("minimal", "low", "medium", "high") -def test_model_group_info_reasoning_efforts_ignore_a_deployment_off_the_map(): +def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_off_the_map(): """The router fills every ModelInfo key, so a deployment absent from the model map arrives with supports_reasoning None rather than with the key missing. Its synthesized entry carries no mode, which is what separates it from a mapped non-reasoning model, and nothing being known about it is - no reason to drop the levels the rest of the group agrees on.""" + no evidence that the unknown deployment accepts levels its mapped sibling supports. The group + therefore reports unknown instead of advertising a value routing might send to either one.""" router = litellm.Router( model_list=[ { @@ -9821,7 +9823,7 @@ def test_model_group_info_reasoning_efforts_ignore_a_deployment_off_the_map(): ) assert result is not None - assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high", "max") + assert result.supported_reasoning_efforts is None @@ -9958,11 +9960,11 @@ def test_model_group_info_survives_a_junk_typed_operator_effort_value(): assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") -def test_model_group_info_reasoning_efforts_ignore_a_mode_the_operator_declared(): +def test_model_group_info_reasoning_efforts_are_unknown_for_an_operator_declared_mode(): """A deployment is registered in the cost map under its own id with whatever model_info the operator wrote, so a mode they set themselves reads back exactly like one the map supplied. Only - a mode the map supplied marks the deployment as known, or an off-map deployment carrying any - mode empties the group it sits in.""" + a mode the map supplied marks the deployment as known. An off-map deployment carrying an + operator mode remains unknown and must keep the whole group's level support unknown.""" from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts mapped_model = "openai/gpt-5.6-sol" @@ -9993,7 +9995,7 @@ def test_model_group_info_reasoning_efforts_ignore_a_mode_the_operator_declared( ) assert result is not None - assert result.supported_reasoning_efforts == expected + assert result.supported_reasoning_efforts is None class TestAddDeploymentApiBaseProviderResolution: @@ -12252,6 +12254,120 @@ class TestTierParamsTheTargetAccepts: assert accepted == {"reasoning_effort": "max"} +class TestRequestReasoningEffortOverride: + def test_drop_effort_from_nested_carrier_preserves_other_nested_values(self): + params: dict[str, object] = {"output_config": {"effort": "high", "format": "json"}} + + litellm.Router._pop_effort_from_nested_carrier(params, "output_config") + + assert params == {"output_config": {"format": "json"}} + + @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) + def test_is_classifier_internal_call_recognizes_both_metadata_carriers(self, metadata_key): + kwargs = {metadata_key: {"internal_call_origin": "autorouter_classifier"}} + + assert litellm.Router._is_classifier_internal_call(kwargs) is True + assert litellm.Router._is_classifier_internal_call({metadata_key: {}}) is False + + def test_removes_every_deployment_native_effort_carrier_without_mutating_shared_config(self): + extra_body: dict[str, object] = { + "reasoning_effort": "high", + "thinking": {"type": "enabled"}, + "output_config": {"effort": "high", "format": "json"}, + "reasoning": {"effort": "high", "summary": "detailed"}, + "provider_option": True, + } + deployment_params: dict[str, object] = { + "model": "bedrock/converse/anthropic.claude-3-7-sonnet", + "thinking": {"type": "enabled", "budget_tokens": 2048}, + "output_config": {"effort": "high", "format": {"type": "json_schema"}}, + "reasoning": {"effort": "high", "summary": "auto"}, + "extra_body": extra_body, + } + + sanitized = litellm.Router._deployment_params_with_request_reasoning_override( + deployment_params, {"reasoning_effort": "low"} + ) + + assert sanitized == { + "model": "bedrock/converse/anthropic.claude-3-7-sonnet", + "output_config": {"format": {"type": "json_schema"}}, + "reasoning": {"summary": "auto"}, + "extra_body": { + "output_config": {"format": "json"}, + "reasoning": {"summary": "detailed"}, + "provider_option": True, + }, + } + assert deployment_params["thinking"] == {"type": "enabled", "budget_tokens": 2048} + assert deployment_params["output_config"] == {"effort": "high", "format": {"type": "json_schema"}} + assert extra_body["reasoning_effort"] == "high" + + @pytest.mark.parametrize("request_kwargs", [{}, {"reasoning_effort": None}]) + def test_omitted_override_preserves_deployment_defaults(self, request_kwargs): + deployment_params = { + "model": "deepseek/deepseek-reasoner", + "thinking": {"type": "enabled"}, + "output_config": {"effort": "high"}, + } + + assert ( + litellm.Router._deployment_params_with_request_reasoning_override(deployment_params, request_kwargs) + == deployment_params + ) + + @pytest.mark.asyncio + async def test_280_concurrent_overrides_never_mutate_or_leak_through_shared_deployment_params(self): + deployment_params = { + "model": "fireworks_ai/accounts/fireworks/models/kimi-k2-thinking", + "thinking": {"type": "enabled"}, + "output_config": {"effort": "high", "format": "json"}, + "extra_body": {"reasoning_effort": "high", "tenant": "shared"}, + } + efforts = ("none", "minimal", "low", "medium", "high", "xhigh", "max") + + results = await asyncio.gather( + *( + asyncio.to_thread( + litellm.Router._deployment_params_with_request_reasoning_override, + deployment_params, + {"reasoning_effort": efforts[index % len(efforts)]}, + ) + for index in range(280) + ) + ) + + assert all("thinking" not in result for result in results) + assert all(result["output_config"] == {"format": "json"} for result in results) + assert all(result["extra_body"] == {"tenant": "shared"} for result in results) + assert deployment_params["thinking"] == {"type": "enabled"} + assert deployment_params["output_config"] == {"effort": "high", "format": "json"} + assert deployment_params["extra_body"] == {"reasoning_effort": "high", "tenant": "shared"} + + @pytest.mark.parametrize( + ("metadata", "should_drop"), + [({"internal_call_origin": "autorouter_classifier"}, True), ({}, False)], + ids=["classifier", "ordinary-request"], + ) + def test_only_classifier_calls_drop_effort_for_an_unsupported_fallback(self, metadata, should_drop): + router = litellm.Router(model_list=[]) + body: dict[str, object] = {"model": "classifier", "reasoning_effort": "low"} + kwargs: dict[str, object] = { + "reasoning_effort": "low", + "metadata": metadata, + "proxy_server_request": {"body": body}, + } + deployment: DeploymentTypedDict = { + "model_name": "fallback", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + } + + router._drop_unsupported_classifier_reasoning_effort(deployment, "fallback", kwargs) + + assert ("reasoning_effort" not in kwargs) is should_drop + assert ("reasoning_effort" not in body) is should_drop + + class TestPreRoutingTierDrivesFallbacks: """#38832: a complexity/auto router picks a tier behind the router name, but fallback lookup stayed on the router name, so the tier's configured chain never ran and a diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index a29527a20fa..00ee7bd7d6e 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -13,6 +13,8 @@ import ClassifierPromptEditor from "./ClassifierPromptEditor"; import CustomTierPromptEditor from "./CustomTierPromptEditor"; import { RestrictedSection, restrictedBy } from "./TierRestrictions"; import HeuristicScoringConfig from "./HeuristicScoringConfig"; +import ClassifierReasoningEffortSelect from "./ClassifierReasoningEffortSelect"; +import type { ReasoningEffort } from "./complexity_router_tiers"; import { useComplexityScorerDefaults } from "@/app/(dashboard)/hooks/autoRouter/useComplexityScorerDefaults"; import { ClassificationFrequency, @@ -154,6 +156,7 @@ interface ClassificationMethodConfigProps { value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; modelOptions: { value: string; label: string }[]; + effortOptionsByModel: Record; customTechnicalKeywords?: string[]; onCustomTechnicalKeywordsChange?: (keywords: string[]) => void; showValidationErrors?: boolean; @@ -236,6 +239,7 @@ const ClassificationMethodConfig: React.FC = ({ value, onChange, modelOptions, + effortOptionsByModel, customTechnicalKeywords, onCustomTechnicalKeywordsChange, showValidationErrors = false, @@ -251,6 +255,9 @@ const ClassificationMethodConfig: React.FC = ({ const contextBudget = value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS; const contextBudgetQuotesNothing = contextBudget > 0 && contextBudget < MIN_QUOTED_CONTEXT_TURN_CHARS; const classificationRubric = value.classifier_llm_config?.classification_rubric ?? DEFAULT_CLASSIFICATION_RUBRIC; + const classifierModel = value.classifier_llm_config?.model ?? ""; + const classifierReasoningEffort = value.classifier_llm_config?.reasoning_effort; + const explicitlySupportedClassifierEfforts = effortOptionsByModel[classifierModel]; const handleClassifierTypeChange = (classifierType: ClassifierType) => { const nextValue: ComplexityRouterConfigValue = { @@ -299,16 +306,33 @@ const ClassificationMethodConfig: React.FC = ({ }; const handleClassifierModelChange = (model: string) => { + if (model === value.classifier_llm_config?.model) return; + const { reasoning_effort: _reasoningEffort, ...classifierLlmConfig } = value.classifier_llm_config ?? { + model: "", + timeout_ms: DEFAULT_CLASSIFIER_TIMEOUT_MS, + }; onChange({ ...value, classifier_llm_config: { - ...value.classifier_llm_config, + ...classifierLlmConfig, model, - timeout_ms: value.classifier_llm_config?.timeout_ms ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, + timeout_ms: classifierLlmConfig.timeout_ms, }, }); }; + const handleClassifierReasoningEffortChange = (reasoningEffort: ReasoningEffort | undefined) => { + if (!value.classifier_llm_config) return; + const { reasoning_effort: _reasoningEffort, ...classifierLlmConfig } = value.classifier_llm_config; + onChange({ + ...value, + classifier_llm_config: + reasoningEffort === undefined + ? classifierLlmConfig + : { ...classifierLlmConfig, reasoning_effort: reasoningEffort }, + }); + }; + const handleClassifierTimeoutChange = (timeoutMs: number) => { onChange({ ...value, @@ -493,9 +517,16 @@ const ClassificationMethodConfig: React.FC = ({ emptyText="No models found" allowClear={false} className={classifierModelMissing ? "border-destructive" : undefined} + aria-label="Classifier Model" /> {classifierModelMissing && A classifier model is required} +