From 26805f75fe44170f36b54d4e8e57d9ce7ad19b5b Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 5 Aug 2026 19:36:37 +0000 Subject: [PATCH 01/67] 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/67] 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 429ad06972b7263435a551a1a5f12f3f985afb53 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:23:57 -0700 Subject: [PATCH 03/67] fix(guardrails): write structured_messages rewrites back into /v1/responses input Message-rewriting guardrails such as Headroom return their rewrite in structured_messages and leave texts untouched. The responses guardrail translation only mapped texts back, so compression never reached the upstream request on /v1/responses while the retrieve tool still got injected. Convert the returned messages back to Responses input (plus instructions) the way the chat and Anthropic handlers already do, and keep developer messages as input_text in the chat-to-responses bridge. Resolves LIT-6494 --- .../transformation.py | 2 +- .../guardrail_translation/handler.py | 47 +++++++-- ...responses_transformation_transformation.py | 13 +++ ...test_openai_responses_guardrail_handler.py | 96 +++++++++++++++++++ .../guardrail_hooks/test_headroom.py | 34 +++++++ 5 files changed, 181 insertions(+), 11 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 85fb0bc8dc6..63c66a4c2c5 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -931,7 +931,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode) def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]: - if role == "user" or role == "system" or role == "tool": + if role in ("user", "system", "developer", "tool"): return {"type": "input_text", "text": content} else: return {"type": "output_text", "text": content} diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7c5d8ac99ad..4320edb8414 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -38,6 +38,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, OpenAiResponsesToChatCompletionStreamIterator, ) from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation @@ -149,8 +150,15 @@ class OpenAIResponsesHandler(BaseTranslation): input_type="request", logging_obj=litellm_logging_obj, ) - guardrailed_texts = guardrailed_inputs.get("texts", []) - data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data + guardrailed_structured_messages = guardrailed_inputs.get("structured_messages") + if ( + guardrailed_structured_messages is not None + and guardrailed_structured_messages is not structured_messages + ): + self._write_back_structured_messages(data, guardrailed_structured_messages) + else: + guardrailed_texts = guardrailed_inputs.get("texts", []) + data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools")) verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") return data @@ -198,24 +206,43 @@ class OpenAIResponsesHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts = guardrailed_inputs.get("texts", []) self._apply_guardrailed_tools_to_data( data, original_tools_list, guardrailed_inputs.get("tools"), ) - # Step 3: Map guardrail responses back to original input structure - await self._apply_guardrail_responses_to_input( - messages=input_data, - responses=guardrailed_texts, - task_mappings=task_mappings, - ) + guardrailed_structured_messages = guardrailed_inputs.get("structured_messages") + if ( + guardrailed_structured_messages is not None + and guardrailed_structured_messages is not structured_messages + ): + self._write_back_structured_messages(data, guardrailed_structured_messages) + else: + # Step 3: Map guardrail responses back to original input structure + await self._apply_guardrail_responses_to_input( + messages=input_data, + responses=guardrailed_inputs.get("texts", []), + task_mappings=task_mappings, + ) - verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", input_data) + verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input")) return data + @staticmethod + def _write_back_structured_messages(data: dict, structured_messages: Sequence[AllMessageValues]) -> None: + input_items, instructions = ( + LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( + list(structured_messages) + ) + ) + data["input"] = input_items + if instructions is None: + data.pop("instructions", None) + return + data["instructions"] = instructions + def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from Responses API request (tools[].name for function and custom, tools[].server_label for mcp).""" diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 21b60d7a216..6590718878d 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1295,6 +1295,19 @@ def test_text_plus_tool_calls_sequence(): # ============================================================================= +def test_developer_message_content_uses_input_text(): + handler = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [{"role": "developer", "content": "Always answer in French."}] + ) + + assert instructions is None + assert input_items == [ + {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "Always answer in French."}]} + ] + + def test_tool_message_output_uses_input_text_not_output_text(): """ Test that tool message content uses input_text type, not output_text. diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 447175b09a6..63f8babf760 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1229,3 +1229,99 @@ class TestOpenAIResponsesHandlerToolInjection: names = [t.get("name") for t in result["tools"]] assert "get_weather" in names assert "injected_tool" in names + + +COMPRESSED_MARKER = "[compressed document; retrieve the full text with hash=b573993006976af767214fac]" + + +class StructuredRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites whole messages via structured_messages and leaves + texts untouched, the way message-compressing guardrails do.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_user = next(i for i, m in enumerate(messages) if m.get("role") == "user") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +def _texts(item: dict) -> list[str]: + content = item.get("content") + if isinstance(content, str): + return [content] + return [part["text"] for part in content] + + +class TestStructuredMessagesWriteBack: + """A guardrail's structured_messages rewrite must land in the Responses request, + not only the per-text mapping the chat handler shares with it.""" + + @pytest.mark.asyncio + async def test_list_input_gets_rewritten_messages_and_keeps_instructions(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [ + {"role": "user", "content": "memo " * 400}, + {"role": "assistant", "content": "Understood."}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["instructions"] == "Answer from the memo only." + user_items = [item for item in result["input"] if item.get("role") == "user"] + assert [_texts(item) for item in user_items] == [[COMPRESSED_MARKER], ["What is the codename?"]] + assert not any(item.get("role") == "system" for item in result["input"]) + assert _texts(next(item for item in result["input"] if item.get("role") == "assistant")) == ["Understood."] + + @pytest.mark.asyncio + async def test_string_input_becomes_rewritten_message_list(self): + handler = OpenAIResponsesHandler() + data = {"model": "gpt-5.6", "input": "memo " * 400} + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert [_texts(item) for item in result["input"]] == [[COMPRESSED_MARKER]] + assert "instructions" not in result + + @pytest.mark.asyncio + async def test_developer_item_survives_write_back_as_input_text(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"role": "developer", "content": "Always answer in French."}, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + developer = next(item for item in result["input"] if item.get("role") == "developer") + assert developer["content"] == [{"type": "input_text", "text": "Always answer in French."}] + + @pytest.mark.asyncio + async def test_same_inputs_object_back_keeps_the_text_mapping(self): + handler = OpenAIResponsesHandler() + original_input = [ + {"role": "user", "content": "Hello"}, + {"role": "user", "content": [{"type": "input_text", "text": "Again"}]}, + ] + data = {"model": "gpt-5.6", "input": original_input} + + result = await handler.process_input_messages(data, MockGuardrail()) + + assert result["input"] is original_input + assert [_texts(item) for item in result["input"]] == [["Hello [GUARDRAILED]"], ["Again [GUARDRAILED]"]] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 1fbc975e40a..37b9852d3e1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -954,6 +954,40 @@ async def test_passthrough_handler_does_not_log_headroom_as_run( assert "headroom" not in _applied_guardrails(data) +@pytest.mark.asyncio +async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstream( + guardrail: HeadroomGuardrail, +): + """Regression for LIT-6494: on /v1/responses the compressed messages must be + written back into `input`, not only the retrieve tool into `tools`, or the + model keeps reading the full document and never calls headroom_retrieve.""" + from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler + + data = { + "model": "gpt-5.6", + "instructions": ORIGINAL_MESSAGES[0]["content"], + "input": [{"role": m["role"], "content": m["content"]} for m in ORIGINAL_MESSAGES[1:]], + "tools": [{"type": "function", "name": "get_weather", "parameters": {"type": "object", "properties": {}}}], + } + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH), + ): + result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert result["instructions"] == ORIGINAL_MESSAGES[0]["content"] + assert [item["content"][0]["text"] for item in result["input"]] == [ + COMPRESSED_MESSAGES_WITH_HASH[0]["content"], + ORIGINAL_MESSAGES[2]["content"], + ORIGINAL_MESSAGES[3]["content"], + ] + assert "A" * 5000 not in json.dumps(result["input"]) + assert [tool["name"] for tool in result["tools"]] == ["get_weather", HEADROOM_RETRIEVE_TOOL_NAME] + + @pytest.mark.asyncio async def test_apply_guardrail_http_error_raises(): guardrail = _make_guardrail() From ef96af51211318ea9bbc9d49947c75b71a3f00b5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:34:42 -0700 Subject: [PATCH 04/67] fix(headroom): resolve CCR retrieval on streaming /v1/responses Streaming /v1/responses requests that carried Headroom's retrieve tool were sent upstream as streams, so the model's headroom_retrieve function_call was streamed straight back to a client that never declared the tool and the retrieval never resolved. Chat completions already avoid this by converting the request to non-stream in the pre-call deployment hook, letting the agentic loop resolve the retrieve, and fake-streaming the final answer. The hook now converts responses call types too, the responses handler wraps the resolved result as a fake stream whenever any interception converted the stream (shared converted_stream_requested helper instead of per-integration key checks), and the follow-up request filter drops every non-code-interpreter interception key through is_interception_internal_key. Resolves LIT-6481 --- litellm/llms/custom_httpx/llm_http_handler.py | 11 +- .../guardrail_hooks/headroom/headroom.py | 5 +- litellm/types/integrations/custom_logger.py | 8 ++ .../guardrail_hooks/test_headroom.py | 117 +++++++++++++++++- 4 files changed, 132 insertions(+), 9 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 573ba85416f..52a91a1d06d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -97,9 +97,12 @@ from litellm.types.containers.main import ( ) from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig from litellm.types.integrations.custom_logger import ( + NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, AgenticLoopPlan, AgenticLoopRequestPatch, AgenticLoopSafetyError, + converted_stream_requested, + is_interception_internal_key, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -2932,10 +2935,7 @@ class BaseLLMHTTPHandler: ) result: Final = final_response if final_response is not None else initial_response - interception_converted_stream: Final = litellm_params.get( - "_code_interpreter_interception_converted_stream" - ) or litellm_params.get("_websearch_interception_converted_stream") - if interception_converted_stream and not litellm_params.get("_agentic_loop_depth"): + if converted_stream_requested(litellm_params) and not litellm_params.get("_agentic_loop_depth"): return self._wrap_responses_response_as_fake_stream( result=result, model=model, @@ -5399,8 +5399,7 @@ class BaseLLMHTTPHandler: kwargs_for_followup: Final = { k: v for k, v in kwargs.items() - if not k.startswith("_websearch_interception") - and not k.startswith("_compression_interception") + if not is_interception_internal_key(k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES) and k != "_code_interpreter_interception_converted_stream" and k not in internal_keys and k not in optional_params diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index e2d2fffb2df..08753bdaed9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -49,6 +49,9 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel BYPASS_HEADER: Final = "x-headroom-bypass" +_STREAM_CONVERTIBLE_CALL_TYPES: Final = frozenset( + (CallTypes.completion, CallTypes.acompletion, CallTypes.responses, CallTypes.aresponses) +) HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" _HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})") _HASH_CACHE_TTL_SECONDS: Final = 15 * 60 @@ -724,7 +727,7 @@ class HeadroomGuardrail(CustomGuardrail): ) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type) effective: Final = base_result if base_result is not None else kwargs - if call_type not in (CallTypes.completion, CallTypes.acompletion): + if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES: return base_result if not effective.get("stream"): return base_result diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 9a714e1724e..5de58a20242 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final from pydantic import BaseModel, Field @@ -29,6 +30,13 @@ def is_interception_internal_key( return any(key.startswith(prefix) for prefix in prefixes) +CONVERTED_STREAM_KEYS: Final = frozenset(f"{prefix}_converted_stream" for prefix in INTERCEPTION_INTERNAL_PREFIXES) + + +def converted_stream_requested(params: Mapping[str, object]) -> bool: + return any(bool(params.get(key)) for key in CONVERTED_STREAM_KEYS) + + class AgenticLoopSafetyError(ValueError): """ Raised when an agentic-loop safety rail refuses a rerun. diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 37b9852d3e1..e1c4cc0922a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1984,6 +1984,58 @@ def _openai_text_payload(content: str) -> dict: return _openai_completion_payload({"role": "assistant", "content": content}, "stop") +def _responses_retrieve_tool_definition() -> dict: + return {"type": "function", **_retrieve_tool_definition()["function"]} + + +def _openai_responses_payload(output_item: dict) -> dict: + return { + "id": "resp_ccr", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-4o", + "output": [output_item], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "temperature": 1.0, + "top_p": 1.0, + "text": {"format": {"type": "text"}}, + "truncation": "disabled", + } + + +def _openai_responses_retrieve_call_payload() -> dict: + return _openai_responses_payload( + { + "type": "function_call", + "id": "fc_ccr", + "call_id": "call_ccr", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": CCR_HASH}), + "status": "completed", + } + ) + + +def _openai_responses_text_payload(text: str) -> dict: + return _openai_responses_payload( + { + "type": "message", + "id": "msg_ccr", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ) + + @pytest.mark.parametrize( "call_type, stream, tools, expect_conversion", [ @@ -1992,12 +2044,14 @@ def _openai_text_payload(content: str) -> dict: (CallTypes.acompletion, False, [_retrieve_tool_definition()], False), (CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False), (CallTypes.acompletion, True, None, False), - (CallTypes.aresponses, True, [_retrieve_tool_definition()], False), + (CallTypes.aresponses, True, [_retrieve_tool_definition()], True), + (CallTypes.responses, True, [_responses_retrieve_tool_definition()], True), + (CallTypes.aresponses, False, [_retrieve_tool_definition()], False), (CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False), ], ) @pytest.mark.asyncio -async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions( +async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions_and_responses( guardrail: HeadroomGuardrail, call_type: CallTypes, stream: bool, @@ -2128,6 +2182,65 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( assert not any(key.startswith("_headroom_interception") for key in followup_body) +@pytest.mark.asyncio +async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + """Regression test for LIT-6481: streaming /v1/responses must resolve the + retrieve tool call server-side exactly like streaming /chat/completions does, + instead of streaming a headroom_retrieve function_call to the client.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_text_payload(final_answer)), + ] + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = await litellm.aresponses( + model="openai/gpt-4o", + input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_responses_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + events = [event async for event in response] + + streamed_text = "".join( + getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert streamed_text == final_answer + assert not any("function_call" in str(getattr(event, "type", "")) for event in events) + assert not any( + getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events + ) + mock_get.assert_called_once() + assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + + assert len(upstream.calls) == 2 + followup_body = json.loads(upstream.calls[1].request.content) + assert not followup_body.get("stream") + assert original_content in json.dumps(followup_body["input"]) + assert not any(key.startswith("_headroom_interception") for key in followup_body) + + # --------------------------------------------------------------------------- # LIT-5018: the turn the model is being asked to act on is never compressed. # From 201ada99828bd7e62de465cd4ceed97002af99a6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:41:08 -0700 Subject: [PATCH 05/67] fix(headroom): check the responses converted-stream flag on the agentic kwargs dict --- litellm/llms/custom_httpx/llm_http_handler.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 52a91a1d06d..ad25a44ab7f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2921,6 +2921,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) + agentic_kwargs: Final = dict(litellm_params) final_response: Final = await self._call_agentic_completion_hooks( response=initial_response, model=model, @@ -2930,12 +2931,12 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, stream=False, custom_llm_provider=custom_llm_provider, - kwargs=dict(litellm_params), + kwargs=agentic_kwargs, api_surface="responses", ) result: Final = final_response if final_response is not None else initial_response - if converted_stream_requested(litellm_params) and not litellm_params.get("_agentic_loop_depth"): + if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"): return self._wrap_responses_response_as_fake_stream( result=result, model=model, From 1695b7f7b1931bcc59c50975c8ccdbb460ad7c04 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 15:02:20 -0700 Subject: [PATCH 06/67] fix(headroom): fake-stream converted sync /v1/responses calls too The sync response_api_handler agentic branch returned the completed ResponsesAPIResponse for a request the Headroom guardrail had converted from streaming, so litellm.responses(stream=True) handed callers a non-iterable object. Wrap it in the same fake stream the async path uses --- litellm/llms/custom_httpx/llm_http_handler.py | 14 ++++- .../guardrail_hooks/test_headroom.py | 52 +++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ad25a44ab7f..2ec5377e217 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2744,6 +2744,7 @@ class BaseLLMHTTPHandler: ) if self._has_agentic_completion_hook(logging_obj): + agentic_kwargs: Final = dict(litellm_params) final_response: Final = run_async_function( self._call_agentic_completion_hooks, response=initial_response, @@ -2754,10 +2755,19 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, stream=False, custom_llm_provider=custom_llm_provider, - kwargs=dict(litellm_params), + kwargs=agentic_kwargs, api_surface="responses", ) - return final_response if final_response is not None else initial_response + result: Final = final_response if final_response is not None else initial_response + if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"): + return self._wrap_responses_response_as_fake_stream( + result=result, + model=model, + responses_api_provider_config=responses_api_provider_config, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + ) + return result return initial_response diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index e1c4cc0922a..c297f593617 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -2241,6 +2241,58 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( assert not any(key.startswith("_headroom_interception") for key in followup_body) +def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +): + """The synchronous responses() path converts the stream the same way, so it + must hand back a stream iterator with the resolved answer rather than the + completed response object.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[ + httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_text_payload(final_answer)), + ] + ) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = litellm.responses( + model="openai/gpt-4o", + input=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_responses_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + events = list(response) + + streamed_text = "".join( + getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta" + ) + assert streamed_text == final_answer + assert not any( + getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events + ) + mock_get.assert_called_once() + assert len(upstream.calls) == 2 + assert not json.loads(upstream.calls[1].request.content).get("stream") + + # --------------------------------------------------------------------------- # LIT-5018: the turn the model is being asked to act on is never compressed. # From 6bd3699d4338193b23f8e937061cbab7c4af2e49 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 16:39:14 -0700 Subject: [PATCH 07/67] fix(responses): keep guardrailed input items and bridge stream usage intact - _write_back_structured_messages now patches only the rewritten rows back into the original input items, so reasoning items (encrypted_content), function_call ids, and web_search_call items survive a guardrail rewrite verbatim; rewrites that cannot be row-mapped fall back to the previous full conversion - the responses bridge stream snapshot restores usage hidden in _hidden_params when stream_options is unset, so converted fake streams report real input_tokens instead of 0 --- litellm/llms/custom_httpx/llm_http_handler.py | 4 +- .../guardrail_translation/handler.py | 135 +++++++++++++++++- .../streaming_iterator.py | 6 + ...test_openai_responses_guardrail_handler.py | 129 ++++++++++++++++- .../guardrail_hooks/test_headroom.py | 2 +- .../test_streaming_iterator_transformation.py | 15 ++ 6 files changed, 279 insertions(+), 12 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2ec5377e217..fcb988705eb 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2744,7 +2744,7 @@ class BaseLLMHTTPHandler: ) if self._has_agentic_completion_hook(logging_obj): - agentic_kwargs: Final = dict(litellm_params) + agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place final_response: Final = run_async_function( self._call_agentic_completion_hooks, response=initial_response, @@ -2931,7 +2931,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) - agentic_kwargs: Final = dict(litellm_params) + agentic_kwargs: Final = dict(litellm_params) # mutable-ok: agentic hooks mutate kwargs in place final_response: Final = await self._call_agentic_completion_hooks( response=initial_response, model=model, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 4320edb8414..8c014feeded 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -28,7 +28,8 @@ Output: response.output is List[GenericResponseOutputItem] where each has: - text: str """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall @@ -50,6 +51,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolParam, OpenAIMcpServerTool, + ResponsesAPIOptionalRequestParams, ResponsesAPIStreamEvents, ) from litellm.types.responses.main import ( @@ -81,6 +83,119 @@ class ResponsesStreamChunk(TypedDict, total=False): text: ReadOnly[str] +_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType( + {"function_call_output": "output", "message": "content"} +) + +_EMPTY_RESPONSES_REQUEST: Final[ResponsesAPIOptionalRequestParams] = {} + + +def _item_rewrite_field(item: Mapping[str, object]) -> str | None: + item_type: Final = item.get("type") + if item_type is None: + return "content" if "content" in item else None + if not isinstance(item_type, str): + return None + return _PATCHABLE_ITEM_FIELDS.get(item_type) + + +def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: + field: Final = _item_rewrite_field(item) + if field is None or not isinstance(rewritten, Mapping): + return None + rewritten_content: Final = rewritten.get("content") + if isinstance(item.get(field), str) and isinstance(rewritten_content, str): + return {**item, field: rewritten_content} # mutable-ok: request input items must stay JSON-plain dicts + rewritten_row: Final = cast("AllMessageValues", rewritten) # cast-ok: guardrails hand back chat-shaped rows + converted_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( + [rewritten_row] # mutable-ok: converter signature takes a list + ) + if len(converted_items) != 1 or not isinstance(converted_items[0], Mapping): + return None + first_converted: Final = cast("Mapping[str, object]", converted_items[0]) # cast-ok: isinstance-checked above + converted_value: Final = first_converted.get(field) + if converted_value is None: + return None + return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts + + +def _input_item_provenance( + raw_input: Sequence[object], + expected_messages: Sequence[object], +) -> tuple[Mapping[int, int], frozenset[int]] | None: + if not all(isinstance(item, Mapping) for item in raw_input): + return None + prefixes: Final = tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", raw_input[:count]), # cast-ok: items checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + for count in range(len(raw_input) + 1) + ) + if tuple(prefixes[-1]) != tuple(expected_messages): + return None + item_for_message: Final = MappingProxyType( + { + message_index: item_index + for item_index in range(len(raw_input)) + for message_index in range(len(prefixes[item_index]), len(prefixes[item_index + 1])) + } + ) + tainted: Final = frozenset( + message_index + for item_index in range(len(raw_input)) + for message_index in range(len(prefixes[item_index])) + if prefixes[item_index + 1][message_index] != prefixes[item_index][message_index] + ) + return item_for_message, tainted + + +def _patch_rewritten_rows_into_input( + data: dict, + original_messages: Sequence[object], + structured_messages: Sequence[object], +) -> bool: + raw_input: Final = data.get("input") + if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages): + return False + offset: Final = 1 if data.get("instructions") else 0 + provenance: Final = _input_item_provenance(raw_input, tuple(original_messages)[offset:]) + if provenance is None: + return False + item_for_message, tainted = provenance + changed: Final = tuple( + (index, rewritten) + for index, (original, rewritten) in enumerate(zip(original_messages, structured_messages)) + if original != rewritten + ) + instruction_rewrites: Final = tuple(rewritten for index, rewritten in changed if index < offset) + rewritten_instructions: Final = ( + instruction_rewrites[0].get("content") + if instruction_rewrites and isinstance(instruction_rewrites[0], Mapping) + else None + ) + if instruction_rewrites and not isinstance(rewritten_instructions, str): + return False + body_changes: Final = tuple((index - offset, rewritten) for index, rewritten in changed if index >= offset) + if any(message_index in tainted or message_index not in item_for_message for message_index, _ in body_changes): + return False + replacements: Final = MappingProxyType( + { + item_for_message[message_index]: _rewritten_input_item( + cast("Mapping[str, object]", raw_input[item_for_message[message_index]]), # cast-ok: checked Mappings + rewritten, + ) + for message_index, rewritten in body_changes + } + ) + if len(replacements) != len(body_changes) or any(item is None for item in replacements.values()): + return False + data["input"] = [replacements.get(index, item) for index, item in enumerate(raw_input)] # mutable-ok: JSON body + if isinstance(rewritten_instructions, str): + data["instructions"] = rewritten_instructions + return True + + class OpenAIResponsesHandler(BaseTranslation): """ Handler for processing OpenAI Responses API with guardrails. @@ -155,9 +270,9 @@ class OpenAIResponsesHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not structured_messages ): - self._write_back_structured_messages(data, guardrailed_structured_messages) + self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages) else: - guardrailed_texts = guardrailed_inputs.get("texts", []) + guardrailed_texts = guardrailed_inputs.get("texts") or () data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools")) verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") @@ -217,12 +332,12 @@ class OpenAIResponsesHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not structured_messages ): - self._write_back_structured_messages(data, guardrailed_structured_messages) + self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages) else: # Step 3: Map guardrail responses back to original input structure await self._apply_guardrail_responses_to_input( messages=input_data, - responses=guardrailed_inputs.get("texts", []), + responses=guardrailed_inputs.get("texts", []), # mutable-ok: callee signature takes a list task_mappings=task_mappings, ) @@ -231,10 +346,16 @@ class OpenAIResponsesHandler(BaseTranslation): return data @staticmethod - def _write_back_structured_messages(data: dict, structured_messages: Sequence[AllMessageValues]) -> None: + def _write_back_structured_messages( + data: dict, + original_messages: Sequence[object], + structured_messages: Sequence[AllMessageValues], + ) -> None: + if _patch_rewritten_rows_into_input(data, original_messages, structured_messages): + return input_items, instructions = ( LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( - list(structured_messages) + list(structured_messages) # mutable-ok: converter signature takes a list ) ) data["input"] = input_items diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 8b1eeb30306..f32599e8098 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -557,6 +557,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): hidden_params: Final = getattr(chunk, "_hidden_params", None) if hidden_params is not None: chunk_dict["_hidden_params"] = dict(hidden_params) if isinstance(hidden_params, dict) else hidden_params + if ( + chunk_dict.get("usage") is None + and isinstance(hidden_params, dict) + and hidden_params.get("usage") is not None + ): + chunk_dict["usage"] = hidden_params["usage"] return chunk_dict def create_reasoning_summary_text_done_event( diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 63f8babf760..87449a4b1e8 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1253,6 +1253,43 @@ class StructuredRewriteGuardrail(CustomGuardrail): return {**inputs, "structured_messages": rewritten} +class ToolOutputRewriteGuardrail(CustomGuardrail): + """Guardrail that compresses the first tool-result row, the way Headroom does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_tool = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "tool") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_tool else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +class DroppingRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the first user row and drops the last row, so the + rewrite can only land through the full-conversion fallback.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first_user = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "user") + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if i == first_user else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten[:-1]} + + def _texts(item: dict) -> list[str]: content = item.get("content") if isinstance(content, str): @@ -1296,7 +1333,93 @@ class TestStructuredMessagesWriteBack: assert "instructions" not in result @pytest.mark.asyncio - async def test_developer_item_survives_write_back_as_input_text(self): + async def test_developer_item_preserved_verbatim_by_row_patch(self): + handler = OpenAIResponsesHandler() + developer_item = {"role": "developer", "content": "Always answer in French."} + data = { + "model": "gpt-5.6", + "input": [ + developer_item, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["input"][0] is developer_item + assert developer_item["content"] == "Always answer in French." + assert _texts(result["input"][1]) == [COMPRESSED_MARKER] + assert _texts(result["input"][2]) == ["What is the codename?"] + + @pytest.mark.asyncio + async def test_reasoning_and_function_call_items_survive_tool_output_compression(self): + handler = OpenAIResponsesHandler() + reasoning_item = { + "id": "rs_123", + "type": "reasoning", + "summary": [], + "encrypted_content": "gAAAAA-signed-reasoning", + } + function_call_item = { + "id": "fc_123", + "type": "function_call", + "call_id": "call_abc", + "name": "read_document", + "arguments": '{"path": "memo.txt"}', + "status": "completed", + } + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [ + reasoning_item, + function_call_item, + {"type": "function_call_output", "call_id": "call_abc", "output": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail()) + + assert result["instructions"] == "Answer from the memo only." + assert result["input"][0] is reasoning_item + assert reasoning_item["encrypted_content"] == "gAAAAA-signed-reasoning" + assert result["input"][1] is function_call_item + assert function_call_item["id"] == "fc_123" + assert result["input"][2] == { + "type": "function_call_output", + "call_id": "call_abc", + "output": COMPRESSED_MARKER, + } + assert result["input"][3] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_web_search_call_item_preserved_verbatim(self): + handler = OpenAIResponsesHandler() + web_search_item = { + "id": "ws_123", + "type": "web_search_call", + "status": "completed", + "action": {"type": "search", "query": "codename memo"}, + } + data = { + "model": "gpt-5.6", + "input": [ + web_search_item, + {"role": "user", "content": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert result["input"][0] is web_search_item + assert _texts(result["input"][1]) == [COMPRESSED_MARKER] + assert _texts(result["input"][2]) == ["What is the codename?"] + + @pytest.mark.asyncio + async def test_row_count_change_falls_back_to_full_conversion(self): handler = OpenAIResponsesHandler() data = { "model": "gpt-5.6", @@ -1307,10 +1430,12 @@ class TestStructuredMessagesWriteBack: ], } - result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + result = await handler.process_input_messages(data, DroppingRewriteGuardrail()) + assert len(result["input"]) == 2 developer = next(item for item in result["input"] if item.get("role") == "developer") assert developer["content"] == [{"type": "input_text", "text": "Always answer in French."}] + assert _texts(next(item for item in result["input"] if item.get("role") == "user")) == [COMPRESSED_MARKER] @pytest.mark.asyncio async def test_same_inputs_object_back_keeps_the_text_mapping(self): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index c297f593617..22c4f46e3fb 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -979,7 +979,7 @@ async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstre result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) assert result["instructions"] == ORIGINAL_MESSAGES[0]["content"] - assert [item["content"][0]["text"] for item in result["input"]] == [ + assert [item["content"] for item in result["input"]] == [ COMPRESSED_MESSAGES_WITH_HASH[0]["content"], ORIGINAL_MESSAGES[2]["content"], ORIGINAL_MESSAGES[3]["content"], diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 823f656ddc5..3ffd44b2fbe 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -24,6 +24,7 @@ from litellm.types.utils import ( ModelResponse, ModelResponseStream, StreamingChoices, + Usage, ) CHAT_COMPLETION_ID = "chatcmpl-77d33d09-effa-4cd2-9c0d-c742d4358256" @@ -523,3 +524,17 @@ async def test_streaming_response_id_falls_back_when_upstream_yields_nothing(): assert response_ids assert len(set(response_ids)) == 1 assert response_ids[0].startswith("resp_") + + +def test_completed_event_restores_usage_hidden_by_stream_options_none(): + final_chunk = _chunk("", finish_reason="stop") + final_chunk._hidden_params = {"usage": Usage(prompt_tokens=117, completion_tokens=5, total_tokens=122)} + iterator = _build_iterator([_chunk("the document says hello"), final_chunk]) + + events = list(iterator) + + completed = next( + event for event in events if getattr(event, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + assert completed.response.usage.input_tokens == 117 + assert completed.response.usage.output_tokens == 5 From 98a9a7e525590d9af15d26ea3db0257d14b42d43 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 21:12:41 -0700 Subject: [PATCH 08/67] fix(streaming): carry hidden usage on the async fake-stream final chunk The sync __next__ exhaustion branch stores calculate_total_usage() in the final chunk's _hidden_params when stream_options is None, but the async __anext__ sibling branch never did. Converted (fake) streams, like the ones the Headroom guardrail produces by flipping streaming /v1/responses calls to non-streaming, are consumed async, so their real usage never reached the completion-to-responses bridge and it token-counted from scratch, reporting input_tokens=0. Mirror the sync branch's hidden-usage block into the async exhaustion branch and add a regression test that async-iterates a CustomStreamWrapper over a MockResponseIterator and asserts the final chunk carries the mock response's usage. --- .../litellm_core_utils/streaming_handler.py | 3 ++ .../test_streaming_handler.py | 37 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 1e0b778d244..d9641a6306b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2336,6 +2336,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() + if self.stream_options is None: + usage: Final = calculate_total_usage(chunks=self.chunks) + processed_chunk._hidden_params["usage"] = usage # pyright: ignore[reportPrivateUsage] # sync parity # see sync __next__'s sibling branch: deliberately do NOT restore # here - this chunk is still this call's own data, and restoring # before returning it would corrupt the caller's own log diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 7f54fbfb4c2..7caf401ce38 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4692,3 +4692,40 @@ async def test_async_stream_assembled_response_keeps_vertex_traffic_type(logging assembled = litellm.stream_chunk_builder(chunks=received, messages=[{"role": "user", "content": "hi"}]) assert assembled is not None assert assembled._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND_FLEX" + + +@pytest.mark.asyncio +async def test_async_fake_stream_final_chunk_carries_hidden_usage(logging_obj: Logging): + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + from litellm.types.utils import ModelResponse + + model_response = ModelResponse( + id="chatcmpl-fake-stream", + model="my-random-model", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "hello world"}, + "finish_reason": "stop", + } + ], + ) + model_response.usage = Usage(prompt_tokens=1234, completion_tokens=7, total_tokens=1241) + + wrapper = CustomStreamWrapper( + completion_stream=MockResponseIterator(model_response=model_response), + model="my-random-model", + custom_llm_provider="anthropic", + logging_obj=logging_obj, + ) + + final_chunk = None + async for chunk in wrapper: + final_chunk = chunk + + assert final_chunk is not None + hidden_usage = final_chunk._hidden_params.get("usage") + assert hidden_usage is not None + assert hidden_usage.prompt_tokens == 1234 + assert hidden_usage.completion_tokens == 7 + assert hidden_usage.total_tokens == 1241 From 574010a2cec368ba2fcffbc6fd8717161b33a3b1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 22:04:26 -0700 Subject: [PATCH 09/67] fix(responses): make guardrail input provenance O(n) and guard non-list structured_messages _input_item_provenance converted every input prefix, so an n-item request paid for n+1 full conversions. It now converts each item once, glues consecutive function_call items (plus their trailing-assistant context) into units so the transform's tool_call merging is reproduced inside the unit conversion, and verifies the unit concatenation against one full conversion, bailing to the full-conversion fallback on any mismatch. Messages from multi-item units are tainted, which keeps parallel tool calls patchable exactly like the old prefix pass while unpredicted merges fall back safely. A guardrail handing back a non-list structured_messages payload (the HiddenLayer v2 evaluation dict) previously fell through the length-mismatch fallback and 500ed converting the dict's keys as messages. The write-back is now skipped for non-list payloads, restoring the previous no-write-back behavior on the Responses surface. Also refreshes the compresr texts-mirror docstring, which still claimed the Responses translation cannot round-trip structured_messages. --- .../guardrail_translation/handler.py | 75 ++++- .../guardrail_hooks/compresr/compresr.py | 7 +- ...test_openai_responses_guardrail_handler.py | 296 ++++++++++++++++++ 3 files changed, 365 insertions(+), 13 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 8c014feeded..bd58fb4cdbc 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -29,6 +29,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has: """ from collections.abc import Mapping, Sequence +from itertools import accumulate from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast @@ -119,33 +120,85 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts +def _is_function_call_item(item: object) -> bool: + return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call") + + +def _last_message_role(messages: Sequence[object]) -> str | None: + if not messages: + return None + last: Final = messages[-1] + role: Final = last.get("role") if isinstance(last, Mapping) else getattr(last, "role", None) + return role if isinstance(role, str) else None + + +def _provenance_unit_bounds( + raw_input: Sequence[object], + solo_conversions: Sequence[Sequence[object]], +) -> tuple[tuple[int, int], ...]: + trailing_roles: Final = tuple( + accumulate( + (_last_message_role(messages) for messages in solo_conversions), + lambda previous, current: current if current is not None else previous, + ) + ) + start_indexes: Final = tuple( + index + for index in range(len(raw_input)) + if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") + ) + return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input)))) + + def _input_item_provenance( raw_input: Sequence[object], expected_messages: Sequence[object], ) -> tuple[Mapping[int, int], frozenset[int]] | None: if not all(isinstance(item, Mapping) for item in raw_input): return None - prefixes: Final = tuple( + solo_conversions: Final = tuple( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=cast("ResponseInputParam", raw_input[:count]), # cast-ok: items checked as Mappings above + input=cast("ResponseInputParam", [item]), # cast-ok: items checked as Mappings above responses_api_request=_EMPTY_RESPONSES_REQUEST, ) - for count in range(len(raw_input) + 1) + for item in raw_input ) - if tuple(prefixes[-1]) != tuple(expected_messages): + full_conversion: Final = tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", list(raw_input)), # cast-ok: items checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + ) + if full_conversion != tuple(expected_messages): return None + units: Final = _provenance_unit_bounds(raw_input, solo_conversions) + unit_messages: Final = tuple( + tuple(solo_conversions[start]) + if end - start == 1 + else tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", list(raw_input[start:end])), # cast-ok: checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + ) + for start, end in units + ) + if tuple(message for messages in unit_messages for message in messages) != full_conversion: + return None + boundaries: Final = tuple(accumulate((len(messages) for messages in unit_messages), initial=0)) item_for_message: Final = MappingProxyType( { - message_index: item_index - for item_index in range(len(raw_input)) - for message_index in range(len(prefixes[item_index]), len(prefixes[item_index + 1])) + message_index: start + for unit_index, (start, end) in enumerate(units) + if end - start == 1 + for message_index in range(boundaries[unit_index], boundaries[unit_index + 1]) } ) tainted: Final = frozenset( message_index - for item_index in range(len(raw_input)) - for message_index in range(len(prefixes[item_index])) - if prefixes[item_index + 1][message_index] != prefixes[item_index][message_index] + for unit_index, (start, end) in enumerate(units) + if end - start > 1 + for message_index in range(boundaries[unit_index], boundaries[unit_index + 1]) ) return item_for_message, tainted @@ -351,6 +404,8 @@ class OpenAIResponsesHandler(BaseTranslation): original_messages: Sequence[object], structured_messages: Sequence[AllMessageValues], ) -> None: + if not isinstance(structured_messages, list): + return if _patch_rewritten_rows_into_input(data, original_messages, structured_messages): return input_items, instructions = ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 5c14d03f50e..84b328a9c60 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -916,9 +916,10 @@ class CompresrGuardrail(CustomGuardrail): def _mirror_texts_channel(input_texts: object, applied: _CompressionResult) -> list[object] | None: """Compressed content mirrored into the Responses `texts` channel. - The chat/Anthropic handlers round-trip ``structured_messages``; the - Responses translation cannot rebuild its input from chat messages and - instead writes back through ``texts``. This matches by value, so a + The chat/Anthropic/Responses handlers round-trip + ``structured_messages``; translations without that round-trip write + back through ``texts``, so the compressed content is mirrored there + too. This matches by value, so a replacement is applied only when it is unambiguous: one compression per text, and every occurrence in ``texts`` accounted for by a compressed target. Anything else is left uncompressed rather than risk a wrong or diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 87449a4b1e8..86c023dd7a4 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1450,3 +1450,299 @@ class TestStructuredMessagesWriteBack: assert result["input"] is original_input assert [_texts(item) for item in result["input"]] == [["Hello [GUARDRAILED]"], ["Again [GUARDRAILED]"]] + + +class AllToolOutputsRewriteGuardrail(CustomGuardrail): + """Guardrail that compresses every tool-result row.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if isinstance(m, dict) and m.get("role") == "tool" else m + for m in messages + ] + return {**inputs, "structured_messages": rewritten} + + +class AssistantRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the first assistant row's content.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "assistant") + rewritten = [{**m, "content": COMPRESSED_MARKER} if i == first else m for i, m in enumerate(messages)] + return {**inputs, "structured_messages": rewritten} + + +class DictStructuredMessagesGuardrail(CustomGuardrail): + """Guardrail that hands back a raw evaluation dict instead of a message list, + the way HiddenLayer v2 does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "structured_messages": {"evaluation": "allowed", "messages": []}} + + +def _parallel_tool_call_input() -> list: + return [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"id": "fc_2", "type": "function_call", "call_id": "call_2", "name": "read_b", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo " * 400}, + {"type": "function_call_output", "call_id": "call_2", "output": "note " * 400}, + {"role": "user", "content": "What is the codename?"}, + ] + + +class TestProvenancePatching: + """The O(n) provenance pass must keep patching rewritten rows in place for the + shapes real agent loops produce, and fall back safely everywhere else.""" + + @pytest.mark.asyncio + async def test_parallel_tool_call_outputs_both_patched(self): + handler = OpenAIResponsesHandler() + raw_input = _parallel_tool_call_input() + fc_1, fc_2 = raw_input[0], raw_input[1] + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, AllToolOutputsRewriteGuardrail()) + + assert result["input"][0] is fc_1 + assert result["input"][1] is fc_2 + assert result["input"][2] == {"type": "function_call_output", "call_id": "call_1", "output": COMPRESSED_MARKER} + assert result["input"][3] == {"type": "function_call_output", "call_id": "call_2", "output": COMPRESSED_MARKER} + assert result["input"][4] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_assistant_turn_with_tool_call_keeps_items_verbatim(self): + handler = OpenAIResponsesHandler() + assistant_item = {"role": "assistant", "content": "Let me read the memo."} + function_call_item = { + "id": "fc_9", + "type": "function_call", + "call_id": "call_9", + "name": "read_document", + "arguments": '{"path": "memo.txt"}', + } + data = { + "model": "gpt-5.6", + "input": [ + assistant_item, + function_call_item, + {"type": "function_call_output", "call_id": "call_9", "output": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail()) + + assert result["input"][0] is assistant_item + assert result["input"][1] is function_call_item + assert result["input"][2] == {"type": "function_call_output", "call_id": "call_9", "output": COMPRESSED_MARKER} + + @pytest.mark.asyncio + async def test_rewrite_of_merged_tool_call_message_falls_back(self): + handler = OpenAIResponsesHandler() + raw_input = _parallel_tool_call_input() + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, AssistantRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] + assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + + @pytest.mark.asyncio + async def test_rewrite_of_lone_function_call_message_falls_back(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo memo"}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + raw_input = data["input"] + result = await handler.process_input_messages(data, AssistantRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] + assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + + def test_provenance_bails_on_non_mapping_item(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + + assert _input_item_provenance(["not a mapping"], []) is None + + def test_provenance_bails_when_expected_messages_disagree(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + + assert _input_item_provenance([{"role": "user", "content": "hi"}], [{"role": "user", "content": "bye"}]) is None + + def test_provenance_bails_on_unpredicted_merge(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + raw_input = [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"role": "assistant", "content": "Reading the memo now."}, + ] + expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=raw_input, responses_api_request={} + ) + assert len(expected) == 1 + assert _input_item_provenance(raw_input, expected) is None + + def test_provenance_maps_and_taints_parallel_tool_calls(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + raw_input = _parallel_tool_call_input() + expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=raw_input, responses_api_request={} + ) + provenance = _input_item_provenance(raw_input, expected) + assert provenance is not None + item_for_message, tainted = provenance + assert tainted == {0} + assert dict(item_for_message) == {1: 2, 2: 3, 3: 4} + + +class TestDictStructuredMessagesGuard: + """A guardrail handing back a non-list structured_messages payload must not + blow up the request; the write-back is skipped instead.""" + + @pytest.mark.asyncio + async def test_list_input_survives_dict_structured_messages(self): + handler = OpenAIResponsesHandler() + original_input = [{"role": "user", "content": "Hello"}] + data = {"model": "gpt-5.6", "input": original_input} + + result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail()) + + assert result["input"] is original_input + assert result["input"] == [{"role": "user", "content": "Hello"}] + + @pytest.mark.asyncio + async def test_string_input_survives_dict_structured_messages(self): + handler = OpenAIResponsesHandler() + data = {"model": "gpt-5.6", "input": "Hello there"} + + result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail()) + + assert result["input"] == "Hello there" + + +class SystemRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the system row, the way prompt-hardening guardrails do.""" + + def __init__(self, rewritten_content: Any = COMPRESSED_MARKER): + super().__init__() + self.rewritten_content = rewritten_content + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "system") + rewritten = [ + {**m, "content": self.rewritten_content} if i == first else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +class TestPatchEdgeBranches: + @pytest.mark.asyncio + async def test_multimodal_user_item_rewritten_through_conversion(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"role": "user", "content": [{"type": "input_text", "text": "memo " * 400}]}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert _texts(result["input"][0]) == [COMPRESSED_MARKER] + assert result["input"][1] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_instructions_rewrite_lands_in_instructions_field(self): + handler = OpenAIResponsesHandler() + user_item = {"role": "user", "content": "What is the codename?"} + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [user_item], + } + + result = await handler.process_input_messages(data, SystemRewriteGuardrail()) + + assert result["instructions"] == COMPRESSED_MARKER + assert result["input"][0] is user_item + + @pytest.mark.asyncio + async def test_non_string_instructions_rewrite_falls_back(self): + handler = OpenAIResponsesHandler() + user_item = {"role": "user", "content": "What is the codename?"} + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [user_item], + } + + result = await handler.process_input_messages( + data, SystemRewriteGuardrail(rewritten_content=[{"type": "text", "text": COMPRESSED_MARKER}]) + ) + + assert result["input"][0] is not user_item + + @pytest.mark.asyncio + async def test_unpredicted_merge_falls_back_through_patch(self): + handler = OpenAIResponsesHandler() + raw_input = [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"role": "assistant", "content": "Reading the memo now."}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo memo"}, + {"role": "user", "content": "memo " * 400}, + ] + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + user_items = [item for item in result["input"] if item.get("role") == "user"] + assert _texts(user_items[0]) == [COMPRESSED_MARKER] + + def test_item_rewrite_field_ignores_non_string_type(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _item_rewrite_field + + assert _item_rewrite_field({"type": 123, "content": "hello"}) is None From 3004b12e900dc521e68f3f4e02e2e6bb53d50dd8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:40:40 -0700 Subject: [PATCH 10/67] refactor(responses): unify guardrail input processing and drop mutable request params The staging merge tightened the ruff-strict and type-discipline budgets, so the two `data: dict` parameters the write-back helpers introduced (LIT001) and the 17-branch `process_input_messages` (C901) no longer fit. Fold the duplicated string/list guardrail flow into one path: a pure `_extract_guardrail_inputs` builds the guardrail payload, the write-back helpers become pure functions returning `_RequestFields` (patched input items plus the resulting instructions value), and the request dict is only mutated in `process_input_messages` itself. `_apply_guardrail_responses_to_input` takes Sequence views since it only reads. A non-list `structured_messages` payload now falls through to the plain texts write-back, matching the pre-write-back behavior for guardrails that never touch structured messages. Handler file deltas vs the merge base: LIT001 57 -> 53, LIT002 42 -> 40, LIT010 28 -> 18, C901 3 -> 3. --- .../guardrail_translation/handler.py | 258 ++++++++---------- 1 file changed, 121 insertions(+), 137 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index bd58fb4cdbc..76a755799d0 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -31,7 +31,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has: from collections.abc import Mapping, Sequence from itertools import accumulate from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.tool_param import FunctionToolParam @@ -203,18 +203,28 @@ def _input_item_provenance( return item_for_message, tainted -def _patch_rewritten_rows_into_input( - data: dict, +class _RequestFields(NamedTuple): + input: tuple[object, ...] + instructions: str | None + + +class _ExtractedInputs(NamedTuple): + inputs: GenericGuardrailAPIInputs + task_mappings: tuple[tuple[int, int | None], ...] + + +def _patched_request_fields( + raw_input: object, + instructions: object, original_messages: Sequence[object], structured_messages: Sequence[object], -) -> bool: - raw_input: Final = data.get("input") +) -> _RequestFields | None: if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages): - return False - offset: Final = 1 if data.get("instructions") else 0 + return None + offset: Final = 1 if instructions else 0 provenance: Final = _input_item_provenance(raw_input, tuple(original_messages)[offset:]) if provenance is None: - return False + return None item_for_message, tainted = provenance changed: Final = tuple( (index, rewritten) @@ -225,13 +235,14 @@ def _patch_rewritten_rows_into_input( rewritten_instructions: Final = ( instruction_rewrites[0].get("content") if instruction_rewrites and isinstance(instruction_rewrites[0], Mapping) - else None + else instructions ) - if instruction_rewrites and not isinstance(rewritten_instructions, str): - return False + instructions_value: Final = rewritten_instructions if isinstance(rewritten_instructions, str) else None + if rewritten_instructions is not None and instructions_value is None: + return None body_changes: Final = tuple((index - offset, rewritten) for index, rewritten in changed if index >= offset) if any(message_index in tainted or message_index not in item_for_message for message_index, _ in body_changes): - return False + return None replacements: Final = MappingProxyType( { item_for_message[message_index]: _rewritten_input_item( @@ -242,11 +253,28 @@ def _patch_rewritten_rows_into_input( } ) if len(replacements) != len(body_changes) or any(item is None for item in replacements.values()): - return False - data["input"] = [replacements.get(index, item) for index, item in enumerate(raw_input)] # mutable-ok: JSON body - if isinstance(rewritten_instructions, str): - data["instructions"] = rewritten_instructions - return True + return None + return _RequestFields( + input=tuple(replacements.get(index, item) for index, item in enumerate(raw_input)), + instructions=instructions_value, + ) + + +def _written_back_request_fields( + raw_input: object, + instructions: object, + original_messages: Sequence[object], + structured_messages: Sequence[AllMessageValues], +) -> _RequestFields | None: + if not isinstance(structured_messages, list): + return None + patched: Final = _patched_request_fields(raw_input, instructions, original_messages, structured_messages) + if patched is not None: + return patched + input_items, converted_instructions = ( + LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(structured_messages) + ) + return _RequestFields(input=tuple(input_items), instructions=converted_instructions) class OpenAIResponsesHandler(BaseTranslation): @@ -288,136 +316,92 @@ class OpenAIResponsesHandler(BaseTranslation): Handles both string input and list of message objects. """ input_data: Final[str | ResponseInputParam | None] = data.get("input") - tools_to_check: Final[list[ChatCompletionToolParam]] = [] - if input_data is None: + if not isinstance(input_data, (str, list)): return data - structured_messages: Final = self.get_structured_messages(data) - - # Handle simple string input - if isinstance(input_data, str): - inputs = GenericGuardrailAPIInputs(texts=[input_data]) - original_tools: list[dict[str, object]] = [] - - # Extract and transform tools if present - if "tools" in data and data["tools"]: - original_tools = list(data["tools"]) - self._extract_and_transform_tools(data["tools"], tools_to_check) - if tools_to_check: - inputs["tools"] = tools_to_check - if structured_messages: - inputs["structured_messages"] = structured_messages - # Include model information if available - model = data.get("model") - if model: - inputs["model"] = model - - guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=inputs, - request_data=data, - input_type="request", - logging_obj=litellm_logging_obj, - ) - guardrailed_structured_messages = guardrailed_inputs.get("structured_messages") - if ( - guardrailed_structured_messages is not None - and guardrailed_structured_messages is not structured_messages - ): - self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages) + extracted: Final = self._extract_guardrail_inputs(data, input_data) + if not extracted.inputs.get("texts"): + return data + if structured_messages: + extracted.inputs["structured_messages"] = structured_messages + original_tools: Final[list[dict[str, object]]] = list(data.get("tools") or []) + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=extracted.inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools")) + written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs) + if written_back is not None: + data["input"] = list(written_back.input) # mutable-ok: JSON body + if written_back.instructions is None: + data.pop("instructions", None) else: - guardrailed_texts = guardrailed_inputs.get("texts") or () - data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data - self._apply_guardrailed_tools_to_data(data, original_tools, guardrailed_inputs.get("tools")) - verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") - return data - - # Handle list input (ResponseInputParam) - if not isinstance(input_data, list): - return data + data["instructions"] = written_back.instructions + elif isinstance(input_data, str): + guardrailed_texts: Final = guardrailed_inputs.get("texts") or () + data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data + else: + await self._apply_guardrail_responses_to_input( + messages=input_data, + responses=guardrailed_inputs.get("texts") or (), + task_mappings=extracted.task_mappings, + ) + verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input")) + return data + def _extract_guardrail_inputs( + self, + data: Mapping[str, object], + input_data: "str | ResponseInputParam", + ) -> _ExtractedInputs: texts_to_check: Final[list[str]] = [] images_to_check: Final[list[str]] = [] task_mappings: Final[list[tuple[int, int | None]]] = [] - original_tools_list: Final[list[dict[str, object]]] = list(data.get("tools") or []) - - # Step 1: Extract all text content, images, and tools - for msg_idx, message in enumerate(input_data): - self._extract_input_text_and_images( - message=message, - msg_idx=msg_idx, - texts_to_check=texts_to_check, - images_to_check=images_to_check, - task_mappings=task_mappings, - ) - - # Extract and transform tools if present - if "tools" in data and data["tools"]: - self._extract_and_transform_tools(data["tools"], tools_to_check) - - # Step 2: Apply guardrail to all texts in batch - if texts_to_check: - inputs = GenericGuardrailAPIInputs(texts=texts_to_check) - if images_to_check: - inputs["images"] = images_to_check - if tools_to_check: - inputs["tools"] = tools_to_check - if structured_messages: - inputs["structured_messages"] = structured_messages - # Include model information if available - model = data.get("model") - if model: - inputs["model"] = model - guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=inputs, - request_data=data, - input_type="request", - logging_obj=litellm_logging_obj, - ) - - self._apply_guardrailed_tools_to_data( - data, - original_tools_list, - guardrailed_inputs.get("tools"), - ) - - guardrailed_structured_messages = guardrailed_inputs.get("structured_messages") - if ( - guardrailed_structured_messages is not None - and guardrailed_structured_messages is not structured_messages - ): - self._write_back_structured_messages(data, structured_messages or (), guardrailed_structured_messages) - else: - # Step 3: Map guardrail responses back to original input structure - await self._apply_guardrail_responses_to_input( - messages=input_data, - responses=guardrailed_inputs.get("texts", []), # mutable-ok: callee signature takes a list + tools_to_check: Final[list[ChatCompletionToolParam]] = [] + if isinstance(input_data, str): + texts_to_check.append(input_data) + else: + for msg_idx, message in enumerate(input_data): + self._extract_input_text_and_images( + message=message, + msg_idx=msg_idx, + texts_to_check=texts_to_check, + images_to_check=images_to_check, task_mappings=task_mappings, ) - - verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input")) - - return data + tools: Final = data.get("tools") + if tools: + self._extract_and_transform_tools( + cast("list[FunctionToolParam | OpenAIMcpServerTool]", tools), # cast-ok: request body tools + tools_to_check, + ) + inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + if tools_to_check: + inputs["tools"] = tools_to_check + model: Final = data.get("model") + if isinstance(model, str): + inputs["model"] = model + return _ExtractedInputs(inputs=inputs, task_mappings=tuple(task_mappings)) @staticmethod - def _write_back_structured_messages( - data: dict, - original_messages: Sequence[object], - structured_messages: Sequence[AllMessageValues], - ) -> None: - if not isinstance(structured_messages, list): - return - if _patch_rewritten_rows_into_input(data, original_messages, structured_messages): - return - input_items, instructions = ( - LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( - list(structured_messages) # mutable-ok: converter signature takes a list - ) + def _written_back_request_fields( + data: Mapping[str, object], + structured_messages: Sequence[AllMessageValues] | None, + guardrailed_inputs: GenericGuardrailAPIInputs, + ) -> _RequestFields | None: + guardrailed: Final = guardrailed_inputs.get("structured_messages") + if guardrailed is None or guardrailed is structured_messages: + return None + return _written_back_request_fields( + data.get("input"), + data.get("instructions"), + structured_messages or (), + guardrailed, ) - data["input"] = input_items - if instructions is None: - data.pop("instructions", None) - return - data["instructions"] = instructions def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from Responses API request (tools[].name for function @@ -543,8 +527,8 @@ class OpenAIResponsesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, messages: Any, # Can be List[Dict[str, Any]] or ResponseInputParam - responses: list[str], - task_mappings: list[tuple[int, int | None]], + responses: Sequence[str], + task_mappings: Sequence[tuple[int, int | None]], ) -> None: """ Apply guardrail responses back to input messages. From 58c5223ab4b34d6bf013b7bafb673691e9d23d2a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 30 Aug 2026 12:59:58 -0700 Subject: [PATCH 11/67] refactor(responses): rename write-back helper out of the method's name The recursion detector in code-quality reads the staticmethod _written_back_request_fields calling the module-level function of the same name as a recursive call. Renaming the module-level helper to _patch_or_convert_request_fields removes the shadowing and describes what it does: patch changed rows in place, else fall back to full conversion. --- .../llms/openai/responses/guardrail_translation/handler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index e7a00a15a23..c676eaa7711 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -264,7 +264,7 @@ def _patched_request_fields( ) -def _written_back_request_fields( +def _patch_or_convert_request_fields( raw_input: object, instructions: object, original_messages: Sequence[object], @@ -408,7 +408,7 @@ class OpenAIResponsesHandler(BaseTranslation): guardrailed: Final = guardrailed_inputs.get("structured_messages") if guardrailed is None or guardrailed is structured_messages: return None - return _written_back_request_fields( + return _patch_or_convert_request_fields( data.get("input"), data.get("instructions"), structured_messages or (), From 1072de94de8e9ce6de0e1e6197caa9ab207abbf6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 10:00:55 -0700 Subject: [PATCH 12/67] fix(azure_ai): only reclassify as azure when api_base is a classic Azure OpenAI endpoint --- litellm/llms/azure_ai/chat/transformation.py | 28 +++++++------ .../chat/test_azure_ai_transformation.py | 40 +++++++++++++++++++ .../test_gpt_5_5_model_metadata.py | 3 +- 3 files changed, 56 insertions(+), 15 deletions(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 7fe9d3dec52..8797f4dc400 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -207,20 +207,22 @@ class AzureAIStudioConfig(OpenAIConfig): message["content"] = texts return stripped_messages - def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool: - try: - if "/" in model: - model = model.split("/", 1)[1] - if ( - model in litellm.open_ai_chat_completion_models - or model in litellm.open_ai_text_completion_models - or model in litellm.open_ai_embedding_models - ): - return True - - except Exception: + def _is_foundry_model_inference_base(self, api_base: str) -> bool: + parsed: Final = urlparse(api_base) + host: Final = parsed.hostname + if host is None or not host.endswith(".services.ai.azure.com"): return False - return False + return "/openai/deployments" not in parsed.path + + def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool: + if api_base is None or self._is_foundry_model_inference_base(api_base): + return False + stripped_model: Final = model.split("/", 1)[1] if "/" in model else model + return ( + stripped_model in litellm.open_ai_chat_completion_models + or stripped_model in litellm.open_ai_text_completion_models + or stripped_model in litellm.open_ai_embedding_models + ) def _get_openai_compatible_provider_info( self, diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 11a727c9635..33fbb4e8fc7 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -31,6 +31,46 @@ async def test_get_openai_compatible_provider_info(): assert custom_llm_provider == "azure" +@pytest.mark.parametrize( + "model, api_base, expected_provider", + [ + ("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com", "azure_ai"), + ("azure_ai/gpt-4o", "https://my-resource.services.ai.azure.com/models", "azure_ai"), + ("azure_ai/gpt-5.4-nano", "https://my-resource.services.ai.azure.com", "azure_ai"), + ("azure_ai/gpt-4o", "https://my-resource.openai.azure.com", "azure"), + ( + "azure_ai/gpt-4o", + "https://my-resource.services.ai.azure.com/openai/deployments/gpt-4o/chat/completions" + "?api-version=2024-08-01-preview", + "azure", + ), + ("azure_ai/mistral-large-latest", "https://my-resource.services.ai.azure.com", "azure_ai"), + ("azure_ai/mistral-large-latest", "https://my-resource.openai.azure.com", "azure_ai"), + ], +) +def test_foundry_base_keeps_azure_ai_provider(model: str, api_base: str, expected_provider: str): + """Regression for #38276: a Foundry .services.ai.azure.com base must not be reclassified as azure.""" + config = AzureAIStudioConfig() + ( + _, + _, + custom_llm_provider, + ) = config._get_openai_compatible_provider_info( + model=model, + api_base=api_base, + api_key="my-key", + custom_llm_provider="azure_ai", + ) + assert custom_llm_provider == expected_provider + + +def test_is_azure_openai_model_without_api_base_keeps_azure_ai(): + """Metadata lookups (get_model_info, supports_* checks) carry no api_base and must not flip the provider.""" + config = AzureAIStudioConfig() + assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base=None) is False + assert config._is_azure_openai_model(model="azure_ai/gpt-4o", api_base="https://my-res.openai.azure.com") is True + + def test_azure_ai_validate_environment(): config = AzureAIStudioConfig() headers = config.validate_environment( diff --git a/tests/test_litellm/test_gpt_5_5_model_metadata.py b/tests/test_litellm/test_gpt_5_5_model_metadata.py index 1c12a48ed9d..a60fa9466e6 100644 --- a/tests/test_litellm/test_gpt_5_5_model_metadata.py +++ b/tests/test_litellm/test_gpt_5_5_model_metadata.py @@ -47,8 +47,7 @@ def test_azure_ai_gpt_5_5_model_info(model): routed_model, provider, _, _ = get_llm_provider(model=model) assert routed_model == model.split("/", 1)[1] - # azure_ai/* models resolve under the azure provider in get_llm_provider - assert provider == "azure" + assert provider == "azure_ai" def test_azure_ai_gpt_5_5_backup_matches_main(): From 604f1fde5014d3a966e1a2ee814e28f83bf97549 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:46:33 -0700 Subject: [PATCH 13/67] fix(azure_ai): route Foundry embeddings to the /models inference route --- litellm/llms/azure_ai/chat/transformation.py | 7 +- litellm/llms/azure_ai/common_utils.py | 9 +++ litellm/llms/azure_ai/embed/handler.py | 19 ++++- .../embed/test_azure_ai_embed_handler.py | 69 +++++++++++++++++++ 4 files changed, 96 insertions(+), 8 deletions(-) create mode 100644 tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 8797f4dc400..f2d405e9a17 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM +from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error from litellm.llms.openai.openai import OpenAIConfig @@ -208,11 +209,7 @@ class AzureAIStudioConfig(OpenAIConfig): return stripped_messages def _is_foundry_model_inference_base(self, api_base: str) -> bool: - parsed: Final = urlparse(api_base) - host: Final = parsed.hostname - if host is None or not host.endswith(".services.ai.azure.com"): - return False - return "/openai/deployments" not in parsed.path + return is_foundry_model_inference_base(api_base) def _is_azure_openai_model(self, model: str, api_base: str | None) -> bool: if api_base is None or self._is_foundry_model_inference_base(api_base): diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index 26a90157455..aa34bab5b2e 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,5 +1,6 @@ from collections.abc import Mapping from typing import Final, Literal +from urllib.parse import urlparse import litellm from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter @@ -10,6 +11,14 @@ from litellm.types.router import GenericLiteLLMParams AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"] +def is_foundry_model_inference_base(api_base: str) -> bool: + parsed: Final = urlparse(api_base) + host: Final = parsed.hostname + if host is None or not host.endswith(".services.ai.azure.com"): + return False + return "/openai/deployments" not in parsed.path + + def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None: """ Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment. diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py index 65c3997c099..c65edbf56e6 100644 --- a/litellm/llms/azure_ai/embed/handler.py +++ b/litellm/llms/azure_ai/embed/handler.py @@ -1,8 +1,10 @@ from typing import Final +from urllib.parse import urlsplit, urlunsplit from openai import OpenAI import litellm +from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -16,6 +18,16 @@ from litellm.utils import convert_to_model_response_object from .cohere_transformation import AzureAICohereConfig +def _foundry_models_route_base(api_base: str | None) -> str | None: + if api_base is None or not is_foundry_model_inference_base(api_base): + return api_base + parts: Final = urlsplit(api_base) + path: Final = parts.path.rstrip("/") + if path.endswith("/models"): + return api_base + return urlunsplit((parts.scheme, parts.netloc, f"{path}/models", parts.query, parts.fragment)) + + class AzureAIEmbedding(OpenAIChatCompletion): def _process_response( self, @@ -214,6 +226,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): assemble result in-order, and return """ + resolved_api_base: Final = _foundry_models_route_base(api_base) if aembedding is True: return self.async_embedding( model, @@ -223,7 +236,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): model_response, optional_params, api_key, - api_base, + resolved_api_base, client, ) @@ -245,7 +258,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): model_response=model_response, optional_params=optional_params, api_key=api_key, - api_base=api_base, + api_base=resolved_api_base, client=client, ) @@ -262,7 +275,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): model_response, optional_params, api_key, - api_base, + resolved_api_base, client=(client if client is not None and isinstance(client, OpenAI) else None), aembedding=aembedding, shared_session=shared_session, diff --git a/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py b/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py new file mode 100644 index 00000000000..0401629171d --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/embed/test_azure_ai_embed_handler.py @@ -0,0 +1,69 @@ +import httpx +import pytest +import respx + +from litellm import embedding +from litellm.llms.azure_ai.embed.handler import _foundry_models_route_base + +EMBEDDING_PAYLOAD = { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, +} + + +@pytest.mark.parametrize( + ("api_base", "expected"), + [ + ( + "https://my-foundry.services.ai.azure.com", + "https://my-foundry.services.ai.azure.com/models", + ), + ( + "https://my-foundry.services.ai.azure.com/", + "https://my-foundry.services.ai.azure.com/models", + ), + ( + "https://my-foundry.services.ai.azure.com?api-version=2024-05-01-preview", + "https://my-foundry.services.ai.azure.com/models?api-version=2024-05-01-preview", + ), + ( + "https://my-foundry.services.ai.azure.com/models", + "https://my-foundry.services.ai.azure.com/models", + ), + ( + "https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small", + "https://my-foundry.services.ai.azure.com/openai/deployments/text-embedding-3-small", + ), + ( + "https://my-resource.openai.azure.com", + "https://my-resource.openai.azure.com", + ), + ( + "https://Mistral-serverless.eastus2.models.ai.azure.com", + "https://Mistral-serverless.eastus2.models.ai.azure.com", + ), + (None, None), + ], +) +def test_foundry_models_route_base(api_base, expected): + assert _foundry_models_route_base(api_base) == expected + + +@respx.mock +def test_azure_ai_embedding_calls_foundry_models_route(): + route = respx.post("https://my-foundry.services.ai.azure.com/models/embeddings").mock( + return_value=httpx.Response(200, json=EMBEDDING_PAYLOAD) + ) + + response = embedding( + model="azure_ai/text-embedding-3-small", + input=["hello world"], + api_base="https://my-foundry.services.ai.azure.com", + api_key="fake-key", + ) + + assert route.called + assert response.data is not None + assert len(response.data) == 1 From c493fc855c85041e331f77fc53ad39950089e8f2 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 1 Sep 2026 15:46:14 -0700 Subject: [PATCH 14/67] fix(ui): show MCP servers and agents inherited from access groups on the team overview The team Overview and Settings tabs fed only object_permission into the Object Permissions card, so a team whose access group grants MCP servers or agents read "MCP Servers 0" and "Agents 0" while the Models card next to it already listed the inherited models. /team/info has returned access_group_mcp_server_ids and access_group_agent_ids for a while, nothing in the dashboard read them. ObjectPermissionsView now accepts the inherited ids and MCPServerPermissions / AgentPermissions merge them into their lists with an Inherited tag, deduped against direct grants, so an admin can tell a group grant from a direct one. Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL --- .../components/object_permissions_view.tsx | 12 ++++- .../permissions/AgentPermissions.test.tsx | 50 +++++++++++++++++++ .../permissions/AgentPermissions.tsx | 34 ++++++++++--- .../permissions/MCPServerPermissions.test.tsx | 43 ++++++++++++++++ .../permissions/MCPServerPermissions.tsx | 34 ++++++++++--- .../src/components/team/TeamInfo.test.tsx | 28 +++++++++++ .../src/components/team/TeamInfo.tsx | 10 +++- 7 files changed, 194 insertions(+), 17 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx index 011791fa18d..7d1c4ce4ac9 100644 --- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx +++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx @@ -6,6 +6,8 @@ import type { ObjectPermission } from "./object_permission_types"; interface ObjectPermissionsViewProps { objectPermission?: ObjectPermission | null; + inheritedMcpServerIds?: string[]; + inheritedAgentIds?: string[]; variant?: "card" | "inline"; className?: string; accessToken?: string | null; @@ -13,6 +15,8 @@ interface ObjectPermissionsViewProps { export function ObjectPermissionsView({ objectPermission, + inheritedMcpServerIds = [], + inheritedAgentIds = [], variant = "card", className = "", accessToken, @@ -34,9 +38,15 @@ export function ObjectPermissionsView({ mcpAccessGroups={mcpAccessGroups} mcpToolPermissions={mcpToolPermissions} mcpToolsets={mcpToolsets} + inheritedMcpServers={inheritedMcpServerIds} + accessToken={accessToken} + /> + -

Search tools

{searchTools.length === 0 ? ( diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx new file mode 100644 index 00000000000..14fe6177ac5 --- /dev/null +++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.test.tsx @@ -0,0 +1,50 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { render, screen } from "@testing-library/react"; +import AgentPermissions from "./AgentPermissions"; +import * as networking from "../networking"; + +vi.mock("../networking"); + +describe("AgentPermissions", () => { + const accessToken = "test-token"; + const agentId = "90337622-756e-4f25-98f0-01fc8174aa24"; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("lists agents inherited from access groups with an Inherited tag and counts them", async () => { + vi.mocked(networking.getAgentsList).mockResolvedValue({ + agents: [{ agent_id: agentId, agent_name: "support_agent" }], + }); + + render(); + + expect(await screen.findByText(/support_agent/)).toBeInTheDocument(); + expect(screen.getByText("Inherited")).toBeInTheDocument(); + expect(screen.getByText("1")).toBeInTheDocument(); + expect(screen.queryByText("No agents or access groups configured")).not.toBeInTheDocument(); + expect(networking.getAgentsList).toHaveBeenCalledWith(accessToken); + }); + + it("does not double-list an agent that is both granted directly and inherited", async () => { + vi.mocked(networking.getAgentsList).mockResolvedValue({ + agents: [{ agent_id: agentId, agent_name: "support_agent" }], + }); + + render(); + + expect(await screen.findByText(/support_agent/)).toBeInTheDocument(); + expect(screen.getAllByText(/support_agent/)).toHaveLength(1); + expect(screen.queryByText("Inherited")).not.toBeInTheDocument(); + expect(screen.getByText("1")).toBeInTheDocument(); + }); + + it("shows the empty state when nothing is granted directly or inherited", () => { + render(); + + expect(screen.getByText("No agents or access groups configured")).toBeInTheDocument(); + expect(screen.getByText("0")).toBeInTheDocument(); + expect(networking.getAgentsList).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx index d84bd1834f8..21bacc23ce1 100644 --- a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx @@ -14,16 +14,26 @@ interface Agent { interface AgentPermissionsProps { agents: string[]; agentAccessGroups?: string[]; + inheritedAgents?: string[]; accessToken?: string | null; } -export function AgentPermissions({ agents, agentAccessGroups = [], accessToken }: AgentPermissionsProps) { +const INHERITED_AGENT_TOOLTIP = "Granted through one of the team's access groups"; + +export function AgentPermissions({ + agents, + agentAccessGroups = [], + inheritedAgents = [], + accessToken, +}: AgentPermissionsProps) { const [agentDetails, setAgentDetails] = useState([]); + const inheritedOnlyAgents = inheritedAgents.filter((agent) => !agents.includes(agent)); + const agentIdCount = agents.length + inheritedOnlyAgents.length; // Fetch agent details when component mounts useEffect(() => { const fetchAgentDetails = async () => { - if (accessToken && agents.length > 0) { + if (accessToken && agentIdCount > 0) { try { const response = await getAgentsList(accessToken); if (response && response.agents && Array.isArray(response.agents)) { @@ -35,7 +45,7 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken } } }; fetchAgentDetails(); - }, [accessToken, agents.length]); + }, [accessToken, agentIdCount]); // Function to get display name for agent const getAgentDisplayName = (agentId: string) => { @@ -47,10 +57,11 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken } return agentId; }; - // Merge agents and access groups into one list + // Merge agents, inherited agents and access groups into one list const mergedItems = [ - ...agents.map((agent) => ({ type: "agent", value: agent })), - ...agentAccessGroups.map((group) => ({ type: "accessGroup", value: group })), + ...agents.map((agent) => ({ type: "agent", value: agent, inherited: false })), + ...inheritedOnlyAgents.map((agent) => ({ type: "agent", value: agent, inherited: true })), + ...agentAccessGroups.map((group) => ({ type: "accessGroup", value: group, inherited: false })), ]; const totalCount = mergedItems.length; @@ -76,8 +87,17 @@ export function AgentPermissions({ agents, agentAccessGroups = [], accessToken } {getAgentDisplayName(item.value)} + {item.inherited && ( + + Inherited + + )} - {`Full ID: ${item.value}`} + + {item.inherited + ? `${INHERITED_AGENT_TOOLTIP}. Full ID: ${item.value}` + : `Full ID: ${item.value}`} + ) : ( diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx index c2df945f367..fe7e3a7b971 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx @@ -406,4 +406,47 @@ describe("MCPServerPermissions", () => { ); await waitFor(() => expect(screen.getByText("Blocked")).toHaveAttribute("data-variant", "destructive")); }); + + it("lists servers inherited from access groups with an Inherited tag and counts them", async () => { + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: mockServerId1, server_name: mockServerName1, alias: mockServerName1 }, + ]); + + render( + , + ); + + expect(await screen.findByText(/DW_MCP/)).toBeInTheDocument(); + expect(screen.getByText("Inherited")).toBeInTheDocument(); + expect(screen.getByText("1")).toBeInTheDocument(); + expect(screen.queryByText("No MCP servers, access groups, or toolsets configured")).not.toBeInTheDocument(); + expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken); + }); + + it("does not double-list a server that is both granted directly and inherited", async () => { + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: mockServerId2, server_name: mockServerName2, alias: mockServerName2 }, + ]); + + render( + , + ); + + expect(await screen.findByText(/Test Server/)).toBeInTheDocument(); + expect(screen.getAllByText(/Test Server/)).toHaveLength(1); + expect(screen.queryByText("Inherited")).not.toBeInTheDocument(); + expect(screen.getByText("1")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx index f11bc6b9627..742409f3f8d 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx @@ -11,14 +11,18 @@ interface MCPServerPermissionsProps { mcpAccessGroups?: string[]; mcpToolPermissions?: Record; mcpToolsets?: string[]; + inheritedMcpServers?: string[]; accessToken?: string | null; } +const INHERITED_MCP_SERVER_TOOLTIP = "Granted through one of the team's access groups"; + export function MCPServerPermissions({ mcpServers, mcpAccessGroups = [], mcpToolPermissions = {}, mcpToolsets = [], + inheritedMcpServers = [], accessToken, }: MCPServerPermissionsProps) { const [mcpServerDetails, setMCPServerDetails] = useState([]); @@ -50,10 +54,16 @@ export function MCPServerPermissions({ }); }; + const directServerIds = mcpServers.filter( + (server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL, + ); + const inheritedOnlyServerIds = inheritedMcpServers.filter((server) => !mcpServers.includes(server)); + const serverIdCount = directServerIds.length + inheritedOnlyServerIds.length; + // Fetch MCP server details when component mounts useEffect(() => { const fetchMCPServerDetails = async () => { - if (accessToken && mcpServers.length > 0) { + if (accessToken && serverIdCount > 0) { try { const response = await fetchMCPServers(accessToken); if (response && Array.isArray(response)) { @@ -67,7 +77,7 @@ export function MCPServerPermissions({ } }; fetchMCPServerDetails(); - }, [accessToken, mcpServers.length]); + }, [accessToken, serverIdCount]); // Fetch toolset details useEffect(() => { @@ -98,12 +108,11 @@ export function MCPServerPermissions({ const blocksAllMcpServers = mcpServers.includes(NO_MCP_SERVERS_SENTINEL); const grantsAllProxyMcpServers = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL); - // Merge servers and access groups into one list + // Merge servers, inherited servers and access groups into one list const mergedItems = [ - ...mcpServers - .filter((server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL) - .map((server) => ({ type: "server", value: server })), - ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group })), + ...directServerIds.map((server) => ({ type: "server", value: server, inherited: false })), + ...inheritedOnlyServerIds.map((server) => ({ type: "server", value: server, inherited: true })), + ...mcpAccessGroups.map((group) => ({ type: "accessGroup", value: group, inherited: false })), ]; const totalCount = mergedItems.length + mcpToolsets.length; @@ -152,8 +161,17 @@ export function MCPServerPermissions({ {getMCPServerDisplayName(item.value)} + {item.inherited && ( + + Inherited + + )} - {`Full ID: ${item.value}`} + + {item.inherited + ? `${INHERITED_MCP_SERVER_TOOLTIP}. Full ID: ${item.value}` + : `Full ID: ${item.value}`} + ) : (
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index ca1e0413dcb..40a3d2ada1a 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -38,6 +38,9 @@ vi.mock("@/components/networking", () => ({ organizationInfoCall: vi.fn(), getRouterSettingsCall: vi.fn().mockResolvedValue({ fields: [] }), getPassThroughEndpointsCall: vi.fn(), + fetchMCPServers: vi.fn().mockResolvedValue([]), + fetchMCPToolsets: vi.fn().mockResolvedValue([]), + getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), })); const can = vi.fn(); @@ -302,6 +305,31 @@ describe("TeamInfoView", () => { ); }); + it("shows MCP servers and agents inherited from access groups in the Object Permissions card", async () => { + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: "mcp-github-1234", server_name: "github", alias: "github" }, + ]); + vi.mocked(networking.getAgentsList).mockResolvedValue({ + agents: [{ agent_id: "agent-support-5678", agent_name: "support_agent" }], + }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + object_permission: null, + access_group_ids: ["ag-1"], + access_group_mcp_server_ids: ["mcp-github-1234"], + access_group_agent_ids: ["agent-support-5678"], + }), + ); + + renderWithProviders(); + + expect(await screen.findByText(/github \(mcp\.\.\.1234\)/)).toBeInTheDocument(); + expect(await screen.findByText(/support_agent \(age\.\.\.5678\)/)).toBeInTheDocument(); + expect(screen.getAllByText("Inherited")).toHaveLength(2); + expect(screen.queryByText("No MCP servers, access groups, or toolsets configured")).not.toBeInTheDocument(); + expect(screen.queryByText("No agents or access groups configured")).not.toBeInTheDocument(); + }); + it("keeps the all-proxy-models badge non-clickable", async () => { vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["all-proxy-models"] })); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 44f0d420fd6..7ebef0691e6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1033,7 +1033,13 @@ const TeamInfoView: React.FC = ({
- + = ({ Date: Tue, 1 Sep 2026 16:25:15 -0700 Subject: [PATCH 15/67] test(ui): query the screen in models page tests and drop restating comments The merge of #38872 into staging kept the destructured render queries in the models-and-endpoints page test, which pushes testing-library/prefer-screen-queries to 21 against a budget of 18 and fails frontend-lint for every PR on top of it. Two comments that only restated the list merge below them are gone as well. Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL --- .../app/(dashboard)/models-and-endpoints/page.test.tsx | 10 +++++----- .../src/components/permissions/AgentPermissions.tsx | 1 - .../components/permissions/MCPServerPermissions.tsx | 1 - 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 1199b66621f..84a05113177 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -111,17 +111,17 @@ describe("ModelsAndEndpointsPage", () => { // POST /model/new 403s a proxy_admin_viewer, so the form's tab must not render for one. it("hides the Add Model tab for a view-only admin session", () => { mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); - const { getByRole, queryByRole } = renderPage(); - expect(queryByRole("tab", { name: "Add Model" })).not.toBeInTheDocument(); - expect(getByRole("tab", { name: "All Models" })).toBeInTheDocument(); + renderPage(); + expect(screen.queryByRole("tab", { name: "Add Model" })).not.toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument(); }); // Read parity: the Auto-Routers list stays reachable for a view-only admin; only the // create affordance inside it is withheld, which AutoRoutersTabPanel decides. it("keeps the Auto-Routers tab for a view-only admin session", () => { mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); - const { getByRole } = renderPage(); - expect(getByRole("tab", { name: /Auto-Routers/ })).toBeInTheDocument(); + renderPage(); + expect(screen.getByRole("tab", { name: /Auto-Routers/ })).toBeInTheDocument(); }); // Auto-routers are excluded from the All Models table, so this tab is their home: the only diff --git a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx index 21bacc23ce1..11f6c414922 100644 --- a/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/AgentPermissions.tsx @@ -57,7 +57,6 @@ export function AgentPermissions({ return agentId; }; - // Merge agents, inherited agents and access groups into one list const mergedItems = [ ...agents.map((agent) => ({ type: "agent", value: agent, inherited: false })), ...inheritedOnlyAgents.map((agent) => ({ type: "agent", value: agent, inherited: true })), diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx index 742409f3f8d..52199dab3c7 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx @@ -108,7 +108,6 @@ export function MCPServerPermissions({ const blocksAllMcpServers = mcpServers.includes(NO_MCP_SERVERS_SENTINEL); const grantsAllProxyMcpServers = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL); - // Merge servers, inherited servers and access groups into one list const mergedItems = [ ...directServerIds.map((server) => ({ type: "server", value: server, inherited: false })), ...inheritedOnlyServerIds.map((server) => ({ type: "server", value: server, inherited: true })), From 7b942fd983af2b033c2767f783777ae86df6cec8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:05:38 -0700 Subject: [PATCH 16/67] fix(azure_ai): route audio and realtime calls on Foundry hosts through the Azure OpenAI handlers --- litellm/constants.py | 1 + litellm/main.py | 10 +++-- litellm/realtime_api/main.py | 3 +- tests/test_litellm/realtime_api/test_main.py | 37 ++++++++++++++++++ tests/test_litellm/test_main.py | 40 ++++++++++++++++++++ 5 files changed, 87 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 1bd977dd9a9..041a83f53e6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_ran DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) +AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"}) ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000 DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) diff --git a/litellm/main.py b/litellm/main.py index c4c5bbefc4f..f6b6453ab45 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -60,6 +60,7 @@ if TYPE_CHECKING: from litellm.types.utils import TokenCountResponse from litellm.constants import ( + AZURE_OPENAI_AUDIO_PROVIDERS, DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, ) @@ -7770,7 +7771,7 @@ def transcription( provider=LlmProviders(custom_llm_provider), ) - if custom_llm_provider == "azure" and provider_config is None: + if custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS and provider_config is None: # azure configs api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") @@ -8057,7 +8058,10 @@ def speech( custom_llm_provider=custom_llm_provider, ) response: HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent] | None = None - if custom_llm_provider == "openai" or custom_llm_provider in litellm.openai_compatible_providers: + if custom_llm_provider == "openai" or ( + custom_llm_provider in litellm.openai_compatible_providers + and custom_llm_provider not in AZURE_OPENAI_AUDIO_PROVIDERS + ): if voice is None or not (isinstance(voice, str)): raise litellm.BadRequestError( message="'voice' is required to be passed as a string for OpenAI TTS", @@ -8111,7 +8115,7 @@ def speech( aspeech=aspeech, shared_session=shared_session, ) - elif custom_llm_provider == "azure": + elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS: # Check if this is Azure Speech Service (Cognitive Services TTS) if model.startswith("speech/"): from litellm.llms.azure.text_to_speech.transformation import ( diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index d4b9f4e8cce..3862aec445f 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -8,6 +8,7 @@ from typing import Any, Final, Literal, cast import litellm from litellm.constants import ( + AZURE_OPENAI_AUDIO_PROVIDERS, REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, request_timeout, @@ -400,7 +401,7 @@ async def _arealtime( litellm_metadata=_build_litellm_metadata(kwargs), query_params=query_params, ) - elif _custom_llm_provider == "azure": + elif _custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS: api_base = dynamic_api_base or litellm_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # set API KEY api_key = dynamic_api_key or litellm.api_key or litellm.openai_key or get_secret_str("AZURE_API_KEY") diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index a3dd5688ad1..761e87ac764 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -1,6 +1,7 @@ import asyncio import time from types import TracebackType +from typing import Final from unittest.mock import MagicMock, patch @@ -294,3 +295,39 @@ async def test_azure_health_check_honors_deployment_realtime_protocol(): model_params={"realtime_protocol": "GA"}, ) assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview" + + +class _ConnectThatStopsAfterCapturingTheUrl: + url: str | None = None + + def __call__(self, url: str, **kwargs: object) -> "_ConnectThatStopsAfterCapturingTheUrl": + self.url = url + return self + + async def __aenter__(self) -> None: + raise RuntimeError("backend url captured, nothing to bridge") + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + return None + + +@pytest.mark.asyncio +async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai_realtime_route(): + connect: Final = _ConnectThatStopsAfterCapturingTheUrl() + with patch("websockets.connect", connect): + await realtime_main._arealtime.__wrapped__( + model="azure_ai/gpt-realtime-mini", + websocket=MagicMock(), + api_base="https://my-project.services.ai.azure.com", + api_key="fake-key", + litellm_logging_obj=FakeLogging(), + ) + assert connect.url == ( + "wss://my-project.services.ai.azure.com/openai/realtime" + "?api-version=2024-10-01-preview&deployment=gpt-realtime-mini" + ) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 8cf878d05d9..30975033b7d 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3181,3 +3181,43 @@ def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent( assert response is not None assert response._hidden_params.get("response_cost") is None + + +FOUNDRY_HOST: Final = "https://my-project.services.ai.azure.com" + + +def test_azure_ai_transcription_on_a_foundry_host_uses_the_azure_openai_deployment_route( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post( + url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/whisper-1/audio/transcriptions\?api-version=.+" + ).mock(return_value=httpx.Response(200, json={"text": "hello"})) + + response: Final = litellm.transcription( + model="azure_ai/whisper-1", + file=("tone.wav", b"RIFF\x00\x00\x00\x00WAVE", "audio/wav"), + api_base=FOUNDRY_HOST, + api_key="fake-key", + ) + + assert route.called + assert response.text == "hello" + + +def test_azure_ai_speech_on_a_foundry_host_uses_the_azure_openai_deployment_route( + respx_mock: respx.MockRouter, +): + route: Final = respx_mock.post( + url__regex=r"https://my-project\.services\.ai\.azure\.com/openai/deployments/tts-1/audio/speech\?api-version=.+" + ).mock(return_value=httpx.Response(200, content=b"mp3-bytes")) + + response: Final = litellm.speech( + model="azure_ai/tts-1", + input="hello", + voice="alloy", + api_base=FOUNDRY_HOST, + api_key="fake-key", + ) + + assert route.called + assert response.content == b"mp3-bytes" 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 17/67] 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 18/67] 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 19/67] 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 de643c028fce8d7634191550968eae7c4c991f04 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 3 Sep 2026 00:25:50 +0000 Subject: [PATCH 20/67] style: format handler after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/openai/responses/guardrail_translation/handler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 0c19ecf909e..59ca601d8be 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -374,7 +374,9 @@ class OpenAIResponsesHandler(BaseTranslation): input_type="request", logging_obj=litellm_logging_obj, ) - self._apply_guardrailed_tools_to_data(data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools")) + self._apply_guardrailed_tools_to_data( + data, original_tools, flattened_tool_groups, guardrailed_inputs.get("tools") + ) written_back: Final = self._written_back_request_fields(data, structured_messages, guardrailed_inputs) if written_back is not None: data["input"] = list(written_back.input) # mutable-ok: JSON body From 45f44fea6019e3e9afd75e7157992899f0846bd0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 17:53:39 -0700 Subject: [PATCH 21/67] fix(headroom): leave background responses requests uncompressed and unconverted --- .../guardrail_hooks/headroom/headroom.py | 6 ++- .../guardrail_hooks/test_headroom.py | 43 +++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 5fbf98d0aae..9d993384461 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -728,6 +728,10 @@ class HeadroomGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER) return inputs + if request_data.get("background"): + verbose_proxy_logger.debug("Headroom: background request; skipping compression") + return inputs + structured_messages: Final = inputs.get("structured_messages") if not _is_object_list(structured_messages) or not structured_messages: return inputs @@ -831,7 +835,7 @@ class HeadroomGuardrail(CustomGuardrail): effective: Final = base_result if base_result is not None else kwargs if call_type not in _STREAM_CONVERTIBLE_CALL_TYPES: return base_result - if not effective.get("stream"): + if not effective.get("stream") or effective.get("background"): return base_result if not has_headroom_retrieve_tool(effective.get("tools")): return base_result diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index a940e4c3394..400eaf8ab3f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -193,6 +193,33 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages( assert "headroom" in _applied_guardrails(request_data) +@pytest.mark.asyncio +async def test_apply_guardrail_leaves_background_requests_uncompressed( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + request_data = {"model": "gpt-4o", "background": True} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES), + ) as post: + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert result is inputs + post.assert_not_awaited() + assert _recorded_guardrail_entries(request_data) == [] + + def _recorded_guardrail_response(request_data: dict) -> dict: entries = request_data["metadata"]["standard_logging_guardrail_information"] assert len(entries) == 1 @@ -2074,6 +2101,22 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple assert kwargs["stream"] is True +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_leaves_background_streams_alone(guardrail: HeadroomGuardrail): + kwargs = { + "model": "gpt-4o", + "stream": True, + "background": True, + "tools": [_responses_retrieve_tool_definition()], + } + + result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.aresponses) + + assert result is kwargs + assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs + assert kwargs["stream"] is True + + @pytest.mark.asyncio async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_configs( guardrail: HeadroomGuardrail, 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 22/67] 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 23/67] 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 e1b2d9de3c85bdeb8b38dab33fcef28ec42ecac5 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 08:45:23 +0000 Subject: [PATCH 24/67] fix(images): forward gpt-image supported params like background to OpenAI and Azure Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/types/utils.py | 1 + litellm/utils.py | 11 ++++--- tests/test_litellm/test_utils.py | 31 ++++++++++++++++++++ tests/test_litellm/types/test_types_utils.py | 9 ++++++ 4 files changed, 48 insertions(+), 4 deletions(-) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 569fce4f7b8..0bdbb83fbb4 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2543,6 +2543,7 @@ class ImageResponse(OpenAIImageResponse, BaseLiteLLMOpenAIResponseObject): ) super().__init__(created=created, data=_data, usage=_usage) + self.background = kwargs.get("background", None) self.quality = kwargs.get("quality", None) self.output_format = kwargs.get("output_format", None) self.size = kwargs.get("size", None) diff --git a/litellm/utils.py b/litellm/utils.py index ba456fc353b..239673a20d9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3338,6 +3338,9 @@ def get_optional_params_image_gen( continue passed_params[k] = v + provider_supported_params: Final[tuple[str, ...]] = ( + tuple(provider_config.get_supported_openai_params(model=model or "")) if provider_config is not None else () + ) default_params: Final = { "n": None, "quality": None, @@ -3348,6 +3351,7 @@ def get_optional_params_image_gen( "imageConfig": None, "tools": None, "web_search_options": None, + **{k: None for k in provider_supported_params}, } non_default_params: Final = _get_non_default_params( @@ -3407,10 +3411,9 @@ def get_optional_params_image_gen( if size is not None: optional_params["aspectRatio"] = _map_openai_size_to_vertex_ai_aspect_ratio(size) - openai_params: list[str] = list(default_params.keys()) - if provider_config is not None: - supported_params = provider_config.get_supported_openai_params(model=model or "") - openai_params = list(supported_params) + openai_params: Final[list[str]] = ( + list(provider_supported_params) if provider_config is not None else list(default_params.keys()) + ) optional_params = add_provider_specific_params_to_optional_params( optional_params=optional_params, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 200cfd02197..99c69d6bc31 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -330,6 +330,37 @@ def test_get_optional_params_image_gen(): assert optional_params["n"] == 3 +@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure"]) +def test_get_optional_params_image_gen_keeps_gpt_image_supported_params(custom_llm_provider): + """https://github.com/BerriAI/litellm/issues/38649""" + from litellm.types.utils import LlmProviders + + provider_config = ProviderConfigManager.get_provider_image_generation_config( + model="gpt-image-2", provider=LlmProviders(custom_llm_provider) + ) + optional_params = get_optional_params_image_gen( + model="gpt-image-2", + n=1, + size="1024x1024", + custom_llm_provider=custom_llm_provider, + provider_config=provider_config, + background="transparent", + output_format="png", + moderation="low", + output_compression=50, + unknown_param="kept-in-extra-body", + ) + assert optional_params == { + "n": 1, + "size": "1024x1024", + "background": "transparent", + "output_format": "png", + "moderation": "low", + "output_compression": 50, + "extra_body": {"unknown_param": "kept-in-extra-body"}, + } + + def test_get_optional_params_image_gen_vertex_ai_size(): """Test that Vertex AI image generation properly handles size parameter and maps it to aspectRatio""" # Test with various size parameters diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index c081b9e8e0d..554604ab200 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -759,3 +759,12 @@ def test_delta_function_tool_call_unchanged_by_custom_support(): delta = Delta(tool_calls=[{"index": 0, "id": "c2", "type": "function", "function": {"name": "g", "arguments": ""}}]) assert isinstance(delta.tool_calls[0], ChatCompletionDeltaToolCall) assert "custom" not in delta.model_dump()["tool_calls"][0] + + +def test_image_response_keeps_background(): + """https://github.com/BerriAI/litellm/issues/38649""" + from litellm.types.utils import ImageResponse + + response = ImageResponse(created=1, data=[{"b64_json": "aGk="}], background="transparent", output_format="png") + assert response.background == "transparent" + assert response.model_dump()["background"] == "transparent" 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 25/67] 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 26/67] 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 27/67] 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 28/67] 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 29/67] 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 30/67] 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 31/67] 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 32/67] 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 33/67] 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}
+