From 698f608ad6ee78ee7d84d8947c386244ea7b1e4e Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 31 Aug 2026 18:26:15 -0700 Subject: [PATCH 001/107] fix(guardrails): transfer guardrail evaluation metadata to spend logs on success path On successful requests, guardrail evaluations run and store results in request_data["litellm_metadata"]["standard_logging_guardrail_information"], but the spend-log serialization reads from request_data["metadata"], so guardrail evaluations were never reported in Request Logs. The failure path already had this transfer; this adds the same logic to the success path so guardrail evaluation info appears in both success and failure spend logs. Fixes LIT-6314. --- .../proxy/hooks/proxy_track_cost_callback.py | 11 ++ .../hooks/test_proxy_track_cost_callback.py | 132 ++++++++++-------- 2 files changed, 85 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 47aafda2337..ca8341822e9 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -269,6 +269,17 @@ class _ProxyDBLogger(CustomLogger): served_model_id=sl_object.get("model_id") if sl_object is not None else None, router=get_llm_router(), ) + # LIT-6314: post-call guardrail evaluations run after SLP is built, so + # guardrail_information is missing from SLP. Populate it from post-call evals + # before spend-log write so reports show accurate guardrail results. + if sl_object is not None: + guardrail_info_from_hooks: Final = ( + kwargs.get("litellm_metadata", {}).get("standard_logging_guardrail_information") + if isinstance(kwargs.get("litellm_metadata"), dict) + else None + ) + if guardrail_info_from_hooks is not None and not sl_object.get("guardrail_information"): + sl_object["guardrail_information"] = guardrail_info_from_hooks if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 8043a1aca3f..5eb59acf4f5 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,4 +1,3 @@ - from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -70,9 +69,7 @@ async def test_async_post_call_failure_hook(): # Check that metadata was properly updated assert "litellm_params" in call_args["kwargs"] - assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == { - "request_id": "test_request_id" - } + assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {"request_id": "test_request_id"} metadata = call_args["kwargs"]["litellm_params"]["metadata"] assert metadata["user_api_key"] == "test_api_key" assert metadata["status"] == "failure" @@ -336,9 +333,7 @@ async def test_should_continue_failure_tracking_when_budget_release_fails(): ) assert mock_invalidate_budget_reservation_counters.await_count == 1 assert ( - mock_invalidate_budget_reservation_counters.await_args.kwargs[ - "budget_reservation" - ] + mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"] is user_api_key_dict.budget_reservation ) assert user_api_key_dict.budget_reservation["finalized"] is True @@ -433,36 +428,21 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): "entries": [{"counter_key": "spend:key:test_api_key"}], } + assert _get_budget_reservation_from_metadata(metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}) is None assert ( _get_budget_reservation_from_metadata( - metadata={"user_api_key_auth": dict(UserAPIKeyAuth())} - ) - is None - ) - assert ( - _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": UserAPIKeyAuth( - budget_reservation=budget_reservation - ) - } + metadata={"user_api_key_auth": UserAPIKeyAuth(budget_reservation=budget_reservation)} ) == budget_reservation ) assert ( _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": dict( - UserAPIKeyAuth(budget_reservation=budget_reservation) - ) - } + metadata={"user_api_key_auth": dict(UserAPIKeyAuth(budget_reservation=budget_reservation))} ) == budget_reservation ) assert ( - _get_budget_reservation_from_metadata( - metadata={"user_api_key_budget_reservation": budget_reservation} - ) + _get_budget_reservation_from_metadata(metadata={"user_api_key_budget_reservation": budget_reservation}) is budget_reservation ) @@ -470,9 +450,7 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=Exception("db unavailable") - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=Exception("db unavailable")) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -508,9 +486,7 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails(): proxy_logging_obj = MagicMock() db_exception = RuntimeError("db unavailable") - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=db_exception - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -554,12 +530,8 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re budget_reservation=budget_reservation, ) assert mock_log_exception.call_count == 2 - mock_log_exception.assert_any_call( - "Failed to release budget reservation after database update failed" - ) - mock_log_exception.assert_any_call( - "Failed to invalidate budget reservation counters after release failed" - ) + mock_log_exception.assert_any_call("Failed to release budget reservation after database update failed") + mock_log_exception.assert_any_call("Failed to invalidate budget reservation counters after release failed") increment_spend_counters.assert_not_awaited() @@ -1101,10 +1073,7 @@ async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj # standard_logging_object should have been propagated from logging obj assert call_kwargs.get("standard_logging_object") is not None - assert ( - call_kwargs["standard_logging_object"]["trace_id"] - == "trace-id-from-logging-obj" - ) + assert call_kwargs["standard_logging_object"]["trace_id"] == "trace-id-from-logging-obj" # litellm_trace_id should also be propagated as a fallback assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" @@ -1691,9 +1660,7 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): "metadata": {}, "proxy_server_request": {"request_id": "rid"}, "response_cost": 3.5e-05, - "combined_usage_object": Usage( - prompt_tokens=30, completion_tokens=1, total_tokens=31 - ), + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), } with patch( @@ -1772,15 +1739,10 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): assert mock_increment.call_args.kwargs["team_id"] == "team-123" assert mock_increment.call_args.kwargs["org_id"] == "org-456" - update_kwargs = ( - mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs - ) + update_kwargs = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs assert update_kwargs["user_id"] == "mcp-user@example.com" assert update_kwargs["team_id"] == "team-123" - assert ( - kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] - == "mcp-user@example.com" - ) + assert kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com" @pytest.mark.parametrize( @@ -1828,9 +1790,7 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect ], ) @pytest.mark.asyncio -async def test_track_cost_callback_logs_unauthenticated_pass_through_request( - call_type, expect_spend_log -): +async def test_track_cost_callback_logs_unauthenticated_pass_through_request(call_type, expect_spend_log): """Regression for LIT-3782: a pass-through request with auth=false reaches the cost callback with no key/user/team/end-user. Before the fix the spend-log write was skipped and the request never appeared in request/usage logs. It @@ -1876,9 +1836,7 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( end_time=datetime.now(), ) - assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( - 1 if expect_spend_log else 0 - ) + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if expect_spend_log else 0) class _FakeDeploymentLookup: @@ -1989,3 +1947,61 @@ async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_un ) assert charged == ("premium", "tier0") + + +@pytest.mark.asyncio +async def test_proxy_track_cost_callback_carries_guardrail_info_to_sl_object(): + """ + LIT-6314 regression: post-call guardrails append evaluations to + request_data["litellm_metadata"]["standard_logging_guardrail_information"], + but the standard_logging_object (built pre-post-call) lacks this field. + The callback must populate standard_logging_object["guardrail_information"] + from post-call evals before spend-log write so reports show results. + """ + logger = _ProxyDBLogger() + guardrail_info = [ + { + "guardrail_name": "bedrock-guard", + "guardrail_status": "success", + "guardrail_cost": 0.0001, + } + ] + sl_object = {"response_cost": 0.01} + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "litellm_metadata": {"standard_logging_guardrail_information": guardrail_info}, + "litellm_params": { + "metadata": {"user_api_key": "test_key"}, + }, + "call_type": CallTypes.completion.value, + "response_cost": 0.01, + "standard_logging_object": sl_object, + } + + with ( + patch( + "litellm.proxy.proxy_server.increment_spend_counters", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.proxy_server.update_cache", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging, + ): + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_proxy_logging.failed_tracking_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == 1 + assert sl_object.get("guardrail_information") == guardrail_info From ca9e121eb67185debdf1af3dea64ab2e080f3df1 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 31 Aug 2026 18:57:20 -0700 Subject: [PATCH 002/107] fix(guardrails): avoid mutable dict literal in guardrail metadata population --- litellm/proxy/hooks/proxy_track_cost_callback.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index ca8341822e9..b78884f7d95 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -273,13 +273,16 @@ class _ProxyDBLogger(CustomLogger): # guardrail_information is missing from SLP. Populate it from post-call evals # before spend-log write so reports show accurate guardrail results. if sl_object is not None: + litellm_metadata: Final = kwargs.get("litellm_metadata") guardrail_info_from_hooks: Final = ( - kwargs.get("litellm_metadata", {}).get("standard_logging_guardrail_information") - if isinstance(kwargs.get("litellm_metadata"), dict) + litellm_metadata.get("standard_logging_guardrail_information") + if isinstance(litellm_metadata, dict) else None ) if guardrail_info_from_hooks is not None and not sl_object.get("guardrail_information"): - sl_object["guardrail_information"] = guardrail_info_from_hooks + sl_object["guardrail_information"] = ( + guardrail_info_from_hooks # mutable-ok: populate SLP before spend-log write + ) if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) From 686e8ddae7aaac16089de238d09c676b6b62824a Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 31 Aug 2026 19:04:14 -0700 Subject: [PATCH 003/107] test: suppress TQ008 on proxy_server global patches, matching file idiom --- .../proxy/hooks/test_proxy_track_cost_callback.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 5eb59acf4f5..2d5857f95e5 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1980,15 +1980,15 @@ async def test_proxy_track_cost_callback_carries_guardrail_info_to_sl_object(): } with ( - patch( + patch( # test-quality-ok: spend counters are a proxy_server global the callback reads lazily, no seam "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock, ), - patch( + patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock, ), - patch( + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam "litellm.proxy.proxy_server.proxy_logging_obj", ) as mock_proxy_logging, ): From 0e6c4279b6387a26ee43e7ef8b8454bdee65295e Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Mon, 31 Aug 2026 21:35:32 -0700 Subject: [PATCH 004/107] fix(guardrails): record not_run evaluation when scoping leaves nothing to scan Replaces the metadata transfer approach: that block read top-level litellm_metadata which guardrail info never populates, and the SLP builder already reads the nested bucket it lands in, so it was dead code and is reverted. Real cause of missing evaluations: process_input_messages skips apply_guardrail entirely when message scoping (skip_system_message, skip_tool, scan_only_tool_results) leaves no scannable content, so the guardrail shows up in applied_guardrails with no guardrail_information entry. Now records a not_run entry unless the guardrail records its own information. --- .../chat/guardrail_translation/handler.py | 9 ++ .../proxy/hooks/proxy_track_cost_callback.py | 14 -- .../test_openai_guardrail_handler.py | 50 +++++++ .../hooks/test_proxy_track_cost_callback.py | 132 ++++++++---------- 4 files changed, 117 insertions(+), 88 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index de15fefe943..76a7f33454a 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -191,6 +191,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings=tool_call_task_mappings, ) + elif not guardrail_to_apply.records_own_guardrail_information: + # Every guardrail in applied_guardrails needs a persisted evaluation record, + # or request logs report it as silently missing (LIT-6314). + guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response="no scannable content after message scoping", + request_data=data, + guardrail_status="not_run", + ) + verbose_proxy_logger.debug( "OpenAI Chat Completions: Processed input messages: %s", data.get("messages"), diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b78884f7d95..47aafda2337 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -269,20 +269,6 @@ class _ProxyDBLogger(CustomLogger): served_model_id=sl_object.get("model_id") if sl_object is not None else None, router=get_llm_router(), ) - # LIT-6314: post-call guardrail evaluations run after SLP is built, so - # guardrail_information is missing from SLP. Populate it from post-call evals - # before spend-log write so reports show accurate guardrail results. - if sl_object is not None: - litellm_metadata: Final = kwargs.get("litellm_metadata") - guardrail_info_from_hooks: Final = ( - litellm_metadata.get("standard_logging_guardrail_information") - if isinstance(litellm_metadata, dict) - else None - ) - if guardrail_info_from_hooks is not None and not sl_object.get("guardrail_information"): - sl_object["guardrail_information"] = ( - guardrail_info_from_hooks # mutable-ok: populate SLP before spend-log write - ) if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a29e0be4655..24060e94e54 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1559,3 +1559,53 @@ class TestScanOnlyToolResults: assert data["messages"][3]["content"] == "page says [BLOCKED] here" assert data["messages"][3]["tool_call_id"] == "call_1" assert data["messages"][4]["content"] == "and then?" + + +class TestNoScannableContentRecordsNotRun: + """LIT-6314: a guardrail whose scoping leaves nothing to scan must still persist an evaluation record""" + + def _system_only_data(self) -> dict: + return {"messages": [{"role": "system", "content": "SYSTEM-PROMPT"}]} + + def _recorded_entries(self, data: dict) -> list: + metadata = data.get("metadata") or data.get("litellm_metadata") or {} + return metadata.get("standard_logging_guardrail_information") or [] + + @pytest.mark.asyncio + async def test_skipped_scan_records_not_run_entry(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="skip-system-guardrail") + guardrail.skip_system_message_in_guardrail = True + data = self._system_only_data() + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.last_inputs is None, "nothing survived scoping, apply_guardrail must not run" + entries = self._recorded_entries(data) + assert len(entries) == 1 + assert entries[0]["guardrail_name"] == "skip-system-guardrail" + assert entries[0]["guardrail_status"] == "not_run" + + @pytest.mark.asyncio + async def test_self_recording_guardrail_is_left_alone(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="self-recording-guardrail") + guardrail.skip_system_message_in_guardrail = True + guardrail.records_own_guardrail_information = True + data = self._system_only_data() + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.last_inputs is None + assert self._recorded_entries(data) == [] + + @pytest.mark.asyncio + async def test_scannable_content_records_no_extra_entry(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="normal-guardrail") + data = {"messages": [{"role": "user", "content": "hello"}]} + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.last_inputs is not None + assert all(e.get("guardrail_status") != "not_run" for e in self._recorded_entries(data)) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 2d5857f95e5..8043a1aca3f 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,3 +1,4 @@ + from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -69,7 +70,9 @@ async def test_async_post_call_failure_hook(): # Check that metadata was properly updated assert "litellm_params" in call_args["kwargs"] - assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {"request_id": "test_request_id"} + assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == { + "request_id": "test_request_id" + } metadata = call_args["kwargs"]["litellm_params"]["metadata"] assert metadata["user_api_key"] == "test_api_key" assert metadata["status"] == "failure" @@ -333,7 +336,9 @@ async def test_should_continue_failure_tracking_when_budget_release_fails(): ) assert mock_invalidate_budget_reservation_counters.await_count == 1 assert ( - mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"] + mock_invalidate_budget_reservation_counters.await_args.kwargs[ + "budget_reservation" + ] is user_api_key_dict.budget_reservation ) assert user_api_key_dict.budget_reservation["finalized"] is True @@ -428,21 +433,36 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): "entries": [{"counter_key": "spend:key:test_api_key"}], } - assert _get_budget_reservation_from_metadata(metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}) is None assert ( _get_budget_reservation_from_metadata( - metadata={"user_api_key_auth": UserAPIKeyAuth(budget_reservation=budget_reservation)} + metadata={"user_api_key_auth": dict(UserAPIKeyAuth())} + ) + is None + ) + assert ( + _get_budget_reservation_from_metadata( + metadata={ + "user_api_key_auth": UserAPIKeyAuth( + budget_reservation=budget_reservation + ) + } ) == budget_reservation ) assert ( _get_budget_reservation_from_metadata( - metadata={"user_api_key_auth": dict(UserAPIKeyAuth(budget_reservation=budget_reservation))} + metadata={ + "user_api_key_auth": dict( + UserAPIKeyAuth(budget_reservation=budget_reservation) + ) + } ) == budget_reservation ) assert ( - _get_budget_reservation_from_metadata(metadata={"user_api_key_budget_reservation": budget_reservation}) + _get_budget_reservation_from_metadata( + metadata={"user_api_key_budget_reservation": budget_reservation} + ) is budget_reservation ) @@ -450,7 +470,9 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=Exception("db unavailable")) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + side_effect=Exception("db unavailable") + ) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -486,7 +508,9 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails(): proxy_logging_obj = MagicMock() db_exception = RuntimeError("db unavailable") - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + side_effect=db_exception + ) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -530,8 +554,12 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re budget_reservation=budget_reservation, ) assert mock_log_exception.call_count == 2 - mock_log_exception.assert_any_call("Failed to release budget reservation after database update failed") - mock_log_exception.assert_any_call("Failed to invalidate budget reservation counters after release failed") + mock_log_exception.assert_any_call( + "Failed to release budget reservation after database update failed" + ) + mock_log_exception.assert_any_call( + "Failed to invalidate budget reservation counters after release failed" + ) increment_spend_counters.assert_not_awaited() @@ -1073,7 +1101,10 @@ async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj # standard_logging_object should have been propagated from logging obj assert call_kwargs.get("standard_logging_object") is not None - assert call_kwargs["standard_logging_object"]["trace_id"] == "trace-id-from-logging-obj" + assert ( + call_kwargs["standard_logging_object"]["trace_id"] + == "trace-id-from-logging-obj" + ) # litellm_trace_id should also be propagated as a fallback assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" @@ -1660,7 +1691,9 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): "metadata": {}, "proxy_server_request": {"request_id": "rid"}, "response_cost": 3.5e-05, - "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + "combined_usage_object": Usage( + prompt_tokens=30, completion_tokens=1, total_tokens=31 + ), } with patch( @@ -1739,10 +1772,15 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): assert mock_increment.call_args.kwargs["team_id"] == "team-123" assert mock_increment.call_args.kwargs["org_id"] == "org-456" - update_kwargs = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs + update_kwargs = ( + mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs + ) assert update_kwargs["user_id"] == "mcp-user@example.com" assert update_kwargs["team_id"] == "team-123" - assert kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com" + assert ( + kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] + == "mcp-user@example.com" + ) @pytest.mark.parametrize( @@ -1790,7 +1828,9 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect ], ) @pytest.mark.asyncio -async def test_track_cost_callback_logs_unauthenticated_pass_through_request(call_type, expect_spend_log): +async def test_track_cost_callback_logs_unauthenticated_pass_through_request( + call_type, expect_spend_log +): """Regression for LIT-3782: a pass-through request with auth=false reaches the cost callback with no key/user/team/end-user. Before the fix the spend-log write was skipped and the request never appeared in request/usage logs. It @@ -1836,7 +1876,9 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(cal end_time=datetime.now(), ) - assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if expect_spend_log else 0) + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( + 1 if expect_spend_log else 0 + ) class _FakeDeploymentLookup: @@ -1947,61 +1989,3 @@ async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_un ) assert charged == ("premium", "tier0") - - -@pytest.mark.asyncio -async def test_proxy_track_cost_callback_carries_guardrail_info_to_sl_object(): - """ - LIT-6314 regression: post-call guardrails append evaluations to - request_data["litellm_metadata"]["standard_logging_guardrail_information"], - but the standard_logging_object (built pre-post-call) lacks this field. - The callback must populate standard_logging_object["guardrail_information"] - from post-call evals before spend-log write so reports show results. - """ - logger = _ProxyDBLogger() - guardrail_info = [ - { - "guardrail_name": "bedrock-guard", - "guardrail_status": "success", - "guardrail_cost": 0.0001, - } - ] - sl_object = {"response_cost": 0.01} - kwargs = { - "model": "gpt-4", - "messages": [{"role": "user", "content": "Hello"}], - "litellm_metadata": {"standard_logging_guardrail_information": guardrail_info}, - "litellm_params": { - "metadata": {"user_api_key": "test_key"}, - }, - "call_type": CallTypes.completion.value, - "response_cost": 0.01, - "standard_logging_object": sl_object, - } - - with ( - patch( # test-quality-ok: spend counters are a proxy_server global the callback reads lazily, no seam - "litellm.proxy.proxy_server.increment_spend_counters", - new_callable=AsyncMock, - ), - patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam - "litellm.proxy.proxy_server.update_cache", - new_callable=AsyncMock, - ), - patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam - "litellm.proxy.proxy_server.proxy_logging_obj", - ) as mock_proxy_logging, - ): - mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() - mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() - mock_proxy_logging.failed_tracking_alert = AsyncMock() - - await logger._PROXY_track_cost_callback( - kwargs=kwargs, - completion_response=None, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == 1 - assert sl_object.get("guardrail_information") == guardrail_info From 527c36343fd93ebaeffec0d4e07f6f9959f85da9 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Tue, 1 Sep 2026 00:20:14 -0700 Subject: [PATCH 005/107] fix(guardrails): keep not_run entries out of daily evaluation counts --- litellm/proxy/guardrails/usage_tracking.py | 21 +++++++------ .../proxy/guardrails/test_usage_tracking.py | 31 +++++++++++++++++++ 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index b8ae09afc00..ec059acb146 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -316,15 +316,18 @@ async def process_spend_logs_guardrail_usage( guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or "" if not guardrail_id: continue - key = _MetricsKey(guardrail_id, date_key) - daily_guardrail[key]["requests_evaluated"] += 1 - action = _guardrail_status_to_action(entry.get("guardrail_status")) - if action == "passed": - daily_guardrail[key]["passed_count"] += 1 - elif action == "blocked": - daily_guardrail[key]["blocked_count"] += 1 - else: - daily_guardrail[key]["flagged_count"] += 1 + status = entry.get("guardrail_status") + # not_run means the guardrail never evaluated the request: index it for drill-down, keep it out of counts + if status != "not_run": + key = _MetricsKey(guardrail_id, date_key) + daily_guardrail[key]["requests_evaluated"] += 1 + action = _guardrail_status_to_action(status) + if action == "passed": + daily_guardrail[key]["passed_count"] += 1 + elif action == "blocked": + daily_guardrail[key]["blocked_count"] += 1 + else: + daily_guardrail[key]["flagged_count"] += 1 policy_id = entry.get("policy_id") index_rows.append( { diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 6da121703d7..b98b037e7b8 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -306,6 +306,37 @@ async def test_zero_and_non_int_usage_counters_are_skipped(): } +@pytest.mark.asyncio +async def test_not_run_entries_are_indexed_but_not_counted_as_evaluations(): + """ + LIT-6314 records a not_run entry when message scoping leaves a guardrail + nothing to scan. The guardrail never evaluated the request, so counting it + as a passed evaluation would inflate daily pass rates; it still gets an + index row so per-request drill-down finds the spend log. + """ + prisma = _prisma() + logs = [_payload("r1", guardrail_status="not_run"), _payload("r2")] + + await process_spend_logs_guardrail_usage(prisma, logs) + + metrics_create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"] + assert metrics_create["requests_evaluated"] == 1 + assert metrics_create["passed_count"] == 1 + index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"] + assert sorted(row["request_id"] for row in index_rows) == ["r1", "r2"] + + +@pytest.mark.asyncio +async def test_batch_of_only_not_run_entries_writes_no_metrics_row(): + prisma = _prisma() + + await process_spend_logs_guardrail_usage(prisma, [_payload("r1", guardrail_status="not_run")]) + + assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 0 + index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"] + assert [row["request_id"] for row in index_rows] == ["r1"] + + @pytest.mark.asyncio async def test_payload_without_request_id_is_skipped_like_the_metrics_path(): prisma = _prisma() From 45ff658c6f020665bbe997ba6fc93ebf7c4eb2e3 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Tue, 1 Sep 2026 01:09:01 -0700 Subject: [PATCH 006/107] fix(ui): render not_run guardrail evaluations as not run instead of failed --- ui/litellm-dashboard/eslint-suppressions.json | 2 +- .../GuardrailViewer/GuardrailViewer.test.tsx | 10 +++ .../GuardrailViewer/GuardrailViewer.tsx | 71 +++++++++++++------ 3 files changed, 60 insertions(+), 23 deletions(-) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 7de7373b20b..93966c8d0ee 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2352,7 +2352,7 @@ }, "src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx": { "no-nested-ternary": { - "count": 4 + "count": 3 } }, "src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx": { diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx index b5e04c72440..0fb5504d231 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx @@ -33,6 +33,16 @@ describe("GuardrailViewer", () => { expect(screen.getByText("1235ms")).toBeInTheDocument(); }); + it("renders not_run entries as not run instead of failed", () => { + const data = makeGuardrailInformation({ guardrail_status: "not_run", guardrail_mode: "pre_call" }); + renderWithProviders(); + + expect(screen.getByText(/1 Not run/)).toBeInTheDocument(); + // one NOT RUN badge in the timeline, one in the evaluation card + expect(screen.getAllByText("NOT RUN")).toHaveLength(2); + expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); + }); + it("calculates and displays masked entity totals", async () => { const user = userEvent.setup(); const data = makeGuardrailInformation({ diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx index 863f4117510..08f9002b050 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx @@ -137,6 +137,36 @@ const isEntrySuccess = (entry: GuardrailInformation): boolean => { return (entry.guardrail_status ?? "").toLowerCase() === "success"; }; +const isEntryNotRun = (entry: GuardrailInformation): boolean => { + return (entry.guardrail_status ?? "").toLowerCase() === "not_run"; +}; + +type EntryStatusLabel = "PASSED" | "NOT RUN" | "FAILED"; + +const entryStatusLabel = (entry: GuardrailInformation): EntryStatusLabel => { + if (isEntrySuccess(entry)) return "PASSED"; + if (isEntryNotRun(entry)) return "NOT RUN"; + return "FAILED"; +}; + +const StatusIcon = ({ status }: { status: EntryStatusLabel }) => { + if (status === "PASSED") return ; + if (status === "NOT RUN") return ; + return ; +}; + +const timelineStatusClass = (status: EntryStatusLabel): string => { + if (status === "PASSED") return "bg-success/15 text-success"; + if (status === "NOT RUN") return "bg-muted text-muted-foreground"; + return "bg-destructive/15 text-destructive"; +}; + +const cardStatusClass = (status: EntryStatusLabel): string => { + if (status === "PASSED") return "bg-success/15 text-success border border-success/20"; + if (status === "NOT RUN") return "bg-muted text-muted-foreground border border-border"; + return "bg-destructive/15 text-destructive border border-destructive/20"; +}; + const getRiskColor = (score: number): string => { if (score <= 3) return "text-success bg-success/10 border-success/20"; if (score <= 6) return "text-warning bg-warning/10 border-warning/20"; @@ -318,8 +348,7 @@ interface TimelineEntry { type: "request" | "guardrail" | "llm" | "response"; label: string; offsetMs: number; - status?: string; - isSuccess?: boolean; + status?: EntryStatusLabel; } const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { @@ -348,8 +377,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { type: "guardrail", label: `Pre-call guardrail: ${getDisplayName(e)}`, offsetMs, - status: isEntrySuccess(e) ? "PASSED" : "FAILED", - isSuccess: isEntrySuccess(e), + status: entryStatusLabel(e), }); } @@ -372,8 +400,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { type: "guardrail", label: `During-call guardrail: ${getDisplayName(e)}`, offsetMs, - status: isEntrySuccess(e) ? "PASSED" : "FAILED", - isSuccess: isEntrySuccess(e), + status: entryStatusLabel(e), }); } @@ -384,8 +411,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { type: "guardrail", label: `Post-call guardrail: ${getDisplayName(e)}`, offsetMs, - status: isEntrySuccess(e) ? "PASSED" : "FAILED", - isSuccess: isEntrySuccess(e), + status: entryStatusLabel(e), }); } @@ -410,10 +436,8 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { ) : item.type === "llm" ? ( - ) : item.isSuccess ? ( - ) : ( - + )} {idx < timeline.length - 1 &&
} @@ -427,9 +451,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { {item.status && ( {item.status} @@ -456,6 +478,7 @@ const formatGuardrailCost = (cost: number): string => { const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { const [expanded, setExpanded] = useState(false); const success = isEntrySuccess(entry); + const statusLabel = entryStatusLabel(entry); const totalMasked = getTotalMasked(entry); const displayName = getDisplayName(entry); const durationStr = formatDurationMs(entry.duration); @@ -490,7 +513,9 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { onClick={() => setExpanded(!expanded)} > {/* Status icon */} -
{success ? : }
+
+ +
{/* Name + badges */}
@@ -501,13 +526,9 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { - {success ? "PASSED" : "FAILED"} + {statusLabel} {matchCountStr && ( @@ -673,7 +694,8 @@ const GuardrailViewer = ({ data, accessToken, logEntry }: GuardrailViewerProps) }, [data]); const passedCount = guardrailEntries.filter(isEntrySuccess).length; - const allPassed = passedCount === guardrailEntries.length; + const notRunCount = guardrailEntries.filter(isEntryNotRun).length; + const allPassed = passedCount === guardrailEntries.length - notRunCount; const totalOverheadMs = useMemo(() => { return Math.round(guardrailEntries.reduce((sum, e) => sum + (e.duration ?? 0), 0) * 1000); @@ -728,6 +750,11 @@ const GuardrailViewer = ({ data, accessToken, logEntry }: GuardrailViewerProps) ) : null} {passedCount} Passed + {notRunCount > 0 && ( + + {notRunCount} Not run + + )}
From 26fc1ff221ba445aadf39fa10ead22d759761afe Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Tue, 1 Sep 2026 01:09:01 -0700 Subject: [PATCH 007/107] fix(guardrails): stop readers from scoring not_run evaluations as passed --- litellm/proxy/compliance_checks.py | 3 +- litellm/proxy/guardrails/usage_endpoints.py | 6 +- .../proxy/guardrails/test_usage_endpoints.py | 70 +++++++++++++++++-- .../test_compliance_endpoints.py | 36 +++++++++- .../GuardrailsMonitor/LogViewer.tsx | 11 ++- .../components/GuardrailsMonitor/mockData.ts | 2 +- .../LogDetailContent.test.tsx | 29 +++++++- .../LogDetailsDrawer/LogDetailContent.tsx | 4 +- 8 files changed, 147 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/compliance_checks.py b/litellm/proxy/compliance_checks.py index ff311911742..053c88d10ed 100644 --- a/litellm/proxy/compliance_checks.py +++ b/litellm/proxy/compliance_checks.py @@ -26,7 +26,8 @@ class ComplianceChecker: def __init__(self, data: ComplianceCheckRequest): self.data = data - self.guardrails = data.guardrail_information or [] + # a not_run entry records a guardrail that never evaluated the request, so it cannot evidence compliance + self.guardrails = [g for g in (data.guardrail_information or []) if g.get("guardrail_status") != "not_run"] def _get_guardrails_by_mode(self, mode: str) -> list[dict]: """ diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 7a0edbddca8..75ffc1545dc 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -256,7 +256,7 @@ class UsageDetailResponse(BaseModel): class UsageLogEntry(BaseModel): id: str timestamp: str - action: str # blocked | passed | flagged + action: str # blocked | passed | flagged | not_run score: float | None latency_ms: float | None model: str | None @@ -691,7 +691,9 @@ def _usage_log_entry_from_row( reason_val = None if entry_for_guardrail: st: Final = (entry_for_guardrail.get("guardrail_status") or "").lower() - if "intervened" in st or "block" in st: + if st == "not_run": + action_val = "not_run" + elif "intervened" in st or "block" in st: action_val = "blocked" elif "fail" in st or "error" in st: action_val = "flagged" diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index 1665fa03639..e469aef9a61 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -9,12 +9,10 @@ orphans), and logs missed their logical-name alias. """ from datetime import datetime -from typing import Any, Optional +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest - - from fastapi import HTTPException from prisma.errors import TableNotFoundError @@ -47,7 +45,7 @@ def _yaml_guardrail( guardrail_id: str = "yaml-1", name: str = "yaml-pii", provider: str = "presidio", - info: Optional[dict] = None, + info: dict | None = None, ) -> Guardrail: return Guardrail( guardrail_id=guardrail_id, @@ -433,3 +431,67 @@ async def test_detail_prev_trend_query_is_bounded(): prev_wheres = [w for w in wheres if "lt" in w.get("date", {})] assert prev_wheres assert all("gte" in w["date"] for w in prev_wheres) + + +@pytest.mark.asyncio +async def test_logs_report_not_run_entries_as_not_run_not_passed(): + """LIT-6314: a guardrail that never scanned must not be reported as a pass in the drill-down.""" + index_row = MagicMock() + index_row.request_id = "req-nr" + index_row.guardrail_id = "db-1" + index_row.start_time = datetime(2026, 4, 22) + spend_log = MagicMock() + spend_log.request_id = "req-nr" + spend_log.model = "gpt-4o-mini" + spend_log.startTime = datetime(2026, 4, 22) + spend_log.metadata = { + "guardrail_information": [ + {"guardrail_name": "db-1", "guardrail_status": "not_run", "duration": 0.0}, + ] + } + prisma = _prisma(find_unique=_db_row(), index_find_many=[index_row]) + prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[spend_log]) + handler = _config_handler() + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_logs( + guardrail_id="db-1", + policy_id=None, + page=1, + page_size=50, + action=None, + start_date=START, + end_date=END, + user_api_key_dict=ADMIN, + ) + assert [log.action for log in resp.logs] == ["not_run"] + + +@pytest.mark.asyncio +async def test_logs_action_passed_filter_excludes_not_run_entries(): + """LIT-6314: filtering the drill-down for passes must not return unscanned requests.""" + index_row = MagicMock() + index_row.request_id = "req-nr" + index_row.guardrail_id = "db-1" + index_row.start_time = datetime(2026, 4, 22) + spend_log = MagicMock() + spend_log.request_id = "req-nr" + spend_log.model = "gpt-4o-mini" + spend_log.startTime = datetime(2026, 4, 22) + spend_log.metadata = {"guardrail_information": [{"guardrail_name": "db-1", "guardrail_status": "not_run"}]} + prisma = _prisma(find_unique=_db_row(), index_find_many=[index_row]) + prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[spend_log]) + handler = _config_handler() + p1, p2 = _patches(prisma, handler) + with p1, p2: + resp = await guardrails_usage_logs( + guardrail_id="db-1", + policy_id=None, + page=1, + page_size=50, + action="passed", + start_date=START, + end_date=END, + user_api_key_dict=ADMIN, + ) + assert resp.logs == [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py index dcbe515d5de..8382a5ada96 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py @@ -2,10 +2,8 @@ Unit tests for compliance check endpoints (EU AI Act and GDPR). """ - import pytest - from litellm.proxy.compliance_checks import ComplianceChecker from litellm.types.proxy.compliance_endpoints import ComplianceCheckRequest @@ -591,3 +589,37 @@ class TestModeMatching: continue if matched: assert mode in _guaranteed_modes(g_mode), (g_mode, mode) + + +class TestNotRunGuardrails: + """LIT-6314 logs a not_run entry for a guardrail that message scoping left nothing to scan.""" + + def test_not_run_alone_never_evidences_compliance(self): + data = ComplianceCheckRequest( + request_id="req-601", + user_id="user-1", + model="gpt-4", + timestamp="2026-02-17T00:00:00Z", + guardrail_information=[ + {"guardrail_name": "pii_detection", "guardrail_status": "not_run", "guardrail_mode": "pre_call"}, + ], + ) + results = {c.check_name: c.passed for c in ComplianceChecker(data).check_eu_ai_act()} + assert results["Guardrails applied"] is False + assert results["Content screened before LLM"] is False + assert results["Audit record complete"] is False + + def test_not_run_sibling_does_not_fail_a_passing_request(self): + data = ComplianceCheckRequest( + request_id="req-602", + user_id="user-1", + model="gpt-4", + timestamp="2026-02-17T00:00:00Z", + pii_detected=True, + guardrail_information=[ + {"guardrail_name": "pii_detection", "guardrail_status": "success", "guardrail_mode": "pre_call"}, + {"guardrail_name": "system_only", "guardrail_status": "not_run", "guardrail_mode": "pre_call"}, + ], + ) + results = {c.check_name: c.passed for c in ComplianceChecker(data).check_gdpr()} + assert results["Sensitive data protected"] is True diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx index 8d073feae82..8b4684e87b0 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx @@ -1,4 +1,4 @@ -import { CircleCheck, ChevronDown, TriangleAlert, X } from "lucide-react"; +import { CircleCheck, ChevronDown, MinusCircle, TriangleAlert, X } from "lucide-react"; import { useQuery } from "@tanstack/react-query"; import moment from "moment"; import React, { useState } from "react"; @@ -10,9 +10,16 @@ import type { LogEntry as ViewLogsLogEntry } from "@/components/view_logs/column import type { LogEntry } from "./mockData"; const actionConfig: Record< - "blocked" | "passed" | "flagged", + "blocked" | "passed" | "flagged" | "not_run", { icon: React.ElementType; color: string; bg: string; border: string; label: string } > = { + not_run: { + icon: MinusCircle, + color: "text-muted-foreground", + bg: "bg-muted", + border: "border-border", + label: "Not run", + }, blocked: { icon: X, color: "text-destructive", diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts b/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts index 7d99ebe7c44..717f0f459e4 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts @@ -43,7 +43,7 @@ export interface LogEntry { input_snippet?: string; output_snippet?: string; score?: number; - action: "blocked" | "passed" | "flagged"; + action: "blocked" | "passed" | "flagged" | "not_run"; model?: string; reason?: string; latency_ms?: number; diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx index a679dc49427..961aad7dc9e 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx @@ -1,7 +1,7 @@ import { render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; -import { LogDetailContent } from "./LogDetailContent"; +import { GuardrailJumpLink, LogDetailContent } from "./LogDetailContent"; import type { LogEntry } from "../columns"; vi.mock("../GuardrailViewer/GuardrailViewer", () => ({ @@ -489,3 +489,30 @@ describe("LogDetailContent", () => { expect(within(descriptions).getByText("-")).toBeInTheDocument(); }); }); + +describe("GuardrailJumpLink", () => { + it("does not render a not_run entry as a failure", () => { + render( + , + ); + expect(screen.getByText(/2 guardrails/)).toHaveTextContent("✓"); + expect(screen.getByText(/2 guardrails/)).not.toHaveTextContent("✗"); + }); + + it("still renders a real failure as failed", () => { + render( + , + ); + expect(screen.getByText(/2 guardrails/)).toHaveTextContent("✗"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx index 4c5c7b7b43f..105ae4add45 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx @@ -636,7 +636,9 @@ function RequestResponseSection({ } export function GuardrailJumpLink({ guardrailEntries }: { guardrailEntries: any[] }) { - const allPassed = guardrailEntries.every((e) => { + // a not_run entry never evaluated the request, so it neither passes nor fails the banner + const evaluated = guardrailEntries.filter((e) => (e?.guardrail_status || e?.status) !== "not_run"); + const allPassed = evaluated.every((e) => { const status = e?.guardrail_status || e?.status; return status === "pass" || status === "passed" || status === "success"; }); From d5cf5640b0d2f5451af553024c135e25a9b05c8d Mon Sep 17 00:00:00 2001 From: AaronHowell <111272993+AaronHowell@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:30:29 +0800 Subject: [PATCH 008/107] fix(responses): preserve provider affinity Co-authored-by: Bytechoreographer --- .../encrypted_content_affinity_check.py | 30 +++- .../test_encrypted_content_affinity_check.py | 136 +++++++++++++++++- 2 files changed, 162 insertions(+), 4 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index b623e31ce06..eba2f36b675 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -37,6 +37,7 @@ Safe to enable globally: """ import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, Protocol, cast import httpx @@ -48,6 +49,7 @@ from litellm.exceptions import ( ServiceUnavailableError, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_utils.cooldown_cache import CooldownCacheValue from litellm.types.llms.openai import AllMessageValues @@ -161,11 +163,13 @@ class EncryptedContentAffinityCheck(CustomLogger): @staticmethod def _encryption_boundary_key( litellm_params: object, - ) -> tuple | None: + ) -> tuple[object, object] | None: """ ``(api_base, api_key)`` pair identifying an Azure resource. Two deployments sharing both are interchangeable for ``encrypted_content`` follow-ups; Azure rejects content produced by any other resource. + Missing values are resolved from ``litellm_credential_name`` without + modifying the deployment, and explicit deployment values take precedence. Accepts any object exposing dict-style ``.get(key, default)``: plain dicts (the common case in ``healthy_deployments``) as well as @@ -180,9 +184,29 @@ class EncryptedContentAffinityCheck(CustomLogger): return None api_base: Final = getter("api_base") api_key: Final = getter("api_key") - if not api_base or not api_key: + credential_name: Final = getter("litellm_credential_name") + credential_values: Final[Mapping[str, object] | None] = ( + CredentialAccessor.get_credential_values(credential_name) + if isinstance(credential_name, str) and credential_name and (api_base is None or api_key is None) + else None + ) + effective_api_base: Final = ( + api_base + if api_base is not None + else credential_values.get("api_base") + if credential_values is not None + else None + ) + effective_api_key: Final = ( + api_key + if api_key is not None + else credential_values.get("api_key") + if credential_values is not None + else None + ) + if not effective_api_base or not effective_api_key: return None - return (api_base, api_key) + return (effective_api_base, effective_api_key) def _find_deployments_on_same_encryption_boundary( self, diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index dac991a41c4..46e67983c33 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -21,8 +21,8 @@ from unittest.mock import AsyncMock, patch import pytest - import litellm +from litellm.models.credentials import CredentialItem from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse @@ -1148,6 +1148,140 @@ def test_boundary_key_accepts_pydantic_litellm_params_instance(): ) +def test_boundary_key_resolves_missing_values_from_named_credential(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + with ( + patch.object( # test-quality-ok: credential registry is the direct dependency under test + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], + ) + ): + boundary = EncryptedContentAffinityCheck._encryption_boundary_key({"litellm_credential_name": "account-a"}) + + assert boundary == ("https://account-a.example.com", "credential-key-a") + + +def test_boundary_key_prefers_explicit_values_over_named_credential(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + with ( + patch.object( # test-quality-ok: credential registry is the direct dependency under test + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://credential.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], + ) + ): + boundary = EncryptedContentAffinityCheck._encryption_boundary_key( + { + "api_base": "https://deployment.example.com", + "litellm_credential_name": "account-a", + } + ) + + assert boundary == ("https://deployment.example.com", "credential-key-a") + + +def test_boundary_fallback_matches_deployments_with_same_named_credential_values(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + with ( + patch.object( # test-quality-ok: credential registry is the direct dependency under test + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ), + CredentialItem( + credential_name="account-a-peer", + credential_values={ + "api_base": "https://account-a.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ), + CredentialItem( + credential_name="account-b", + credential_values={ + "api_base": "https://account-b.example.com", + "api_key": "credential-key-b", + }, + credential_info={}, + ), + ], + ) + ): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5.3-codex", + "litellm_params": { + "model": "azure/gpt-5.3-codex", + "litellm_credential_name": "account-a", + }, + "model_info": {"id": "origin"}, + } + ], + num_retries=0, + ) + check = EncryptedContentAffinityCheck(router=router) + healthy_deployments = [ + { + "model_info": {"id": "peer-same-boundary"}, + "litellm_params": { + "model": "azure/gpt-5.4", + "litellm_credential_name": "account-a-peer", + }, + }, + { + "model_info": {"id": "peer-different-boundary"}, + "litellm_params": { + "model": "azure/gpt-5.4", + "litellm_credential_name": "account-b", + }, + }, + ] + + matches, originating = check._find_deployments_on_same_encryption_boundary( + healthy_deployments=healthy_deployments, + model_id="origin", + ) + + assert originating is not None + assert [deployment["model_info"]["id"] for deployment in matches] == ["peer-same-boundary"] + + def test_boundary_key_rejects_non_dict_like_inputs(): """ Inputs that don't expose ``.get()`` (None, lists, strings, ints) -> None. From 2cd28b97f1793ef6032526a0119cb892dc3e9b66 Mon Sep 17 00:00:00 2001 From: AaronHowell <111272993+AaronHowell@users.noreply.github.com> Date: Tue, 8 Sep 2026 18:02:31 +0800 Subject: [PATCH 009/107] fix(responses): align credential boundary resolution --- .../encrypted_content_affinity_check.py | 26 +++++-------- .../test_encrypted_content_affinity_check.py | 37 ++++++++++++++++++- 2 files changed, 45 insertions(+), 18 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index eba2f36b675..db3c865e183 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -165,11 +165,9 @@ class EncryptedContentAffinityCheck(CustomLogger): litellm_params: object, ) -> tuple[object, object] | None: """ - ``(api_base, api_key)`` pair identifying an Azure resource. Two - deployments sharing both are interchangeable for ``encrypted_content`` - follow-ups; Azure rejects content produced by any other resource. - Missing values are resolved from ``litellm_credential_name`` without - modifying the deployment, and explicit deployment values take precedence. + ``(api_base, api_key)`` identifies an upstream encryption boundary. + The values are resolved from the deployment and its named credential + without modifying the deployment. Accepts any object exposing dict-style ``.get(key, default)``: plain dicts (the common case in ``healthy_deployments``) as well as @@ -187,22 +185,18 @@ class EncryptedContentAffinityCheck(CustomLogger): credential_name: Final = getter("litellm_credential_name") credential_values: Final[Mapping[str, object] | None] = ( CredentialAccessor.get_credential_values(credential_name) - if isinstance(credential_name, str) and credential_name and (api_base is None or api_key is None) + if isinstance(credential_name, str) and credential_name else None ) effective_api_base: Final = ( - api_base - if api_base is not None - else credential_values.get("api_base") - if credential_values is not None - else None + credential_values.get("api_base") + if credential_values is not None and "api_base" in credential_values + else api_base ) effective_api_key: Final = ( - api_key - if api_key is not None - else credential_values.get("api_key") - if credential_values is not None - else None + credential_values.get("api_key") + if credential_values is not None and "api_key" in credential_values + else api_key ) if not effective_api_base or not effective_api_key: return None diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 46e67983c33..7b7a4969d41 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -1174,7 +1174,7 @@ def test_boundary_key_resolves_missing_values_from_named_credential(): assert boundary == ("https://account-a.example.com", "credential-key-a") -def test_boundary_key_prefers_explicit_values_over_named_credential(): +def test_boundary_key_matches_named_credential_precedence(): from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) @@ -1198,11 +1198,44 @@ def test_boundary_key_prefers_explicit_values_over_named_credential(): boundary = EncryptedContentAffinityCheck._encryption_boundary_key( { "api_base": "https://deployment.example.com", + "api_key": "deployment-key", "litellm_credential_name": "account-a", } ) - assert boundary == ("https://deployment.example.com", "credential-key-a") + assert boundary == ("https://credential.example.com", "credential-key-a") + + +def test_boundary_key_resolves_credential_when_explicit_values_are_empty(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + with ( + patch.object( # test-quality-ok: credential registry is the direct dependency under test + litellm, + "credential_list", + [ + CredentialItem( + credential_name="account-a", + credential_values={ + "api_base": "https://credential.example.com", + "api_key": "credential-key-a", + }, + credential_info={}, + ) + ], + ) + ): + boundary = EncryptedContentAffinityCheck._encryption_boundary_key( + { + "api_base": "", + "api_key": "", + "litellm_credential_name": "account-a", + } + ) + + assert boundary == ("https://credential.example.com", "credential-key-a") def test_boundary_fallback_matches_deployments_with_same_named_credential_values(): From e2b41286d351c6eda893c7e335d4a03896791f5e Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 00:48:27 -0400 Subject: [PATCH 010/107] fix(vertex-live): bill every modality on the /vertex_ai/live passthrough The Live passthrough builds Usage from the TEXT-modality counts alone, so audio, image and video tokens never reach the cost calculator and bill as nothing. A one-turn audio session reported 13 text and 127 audio input tokens and billed the 13; a camera session reported 1043 prompt tokens and billed 11. Reporting the full per-modality breakdown fixes it, because the shared Gemini input and output cost path already prices audio, image and video from prompt_tokens_details and completion_tokens_details. On the native-audio entry that is a 6x difference per token in both directions, which is the whole gap. Aggregation across turns is unchanged. Google charges per turn for every token in the Live session context window, current turn plus all accumulated tokens from previous turns, so the existing summing is what Vertex bills and it stays as it is. That is worth stating because the cumulative promptTokensDetails looks like a restatement of one running total, and treating it that way would under-bill a multi-turn session. See the LiveAPI context-window note on https://cloud.google.com/vertex-ai/generative-ai/pricing. Live can also name the modality carrying the rest of a turn and omit its tokenCount. Reading that absent key as zero left the tokens inside candidatesTokenCount but outside the breakdown, so real speech was charged at the text output rate. A lone unpriced entry now takes whatever the turn's declared count leaves over. Two or more cannot be told apart, so they are still left to the calculator's text remainder. Server-side toolUsePromptTokenCount is now reported in prompt_tokens_details. It is deliberately kept out of prompt_tokens: no Gemini route prices tool-use tokens, and adding them there instead suppresses the cache-overlap correction and raises the bill for no extra work. Removes _calculate_live_api_cost, whose result never reached the bill. It set kwargs["response_cost"], which the standard logging path recomputes from the ModelResponse, and on a measured audio session it returned $0.000487 against a $0.0000425 row. Now that the modality counts reach the standard calculator, keeping a second hand-rolled pricing path would only ever double-charge. The rewrite of the aggregator is arithmetically identical to what it replaced. It sums the same three counts and the same per-modality details, still takes the remaining fields from the first turn, and drops nine LIT010, one C901 and 42 basedpyright findings in the process. --- ...tex_ai_live_passthrough_logging_handler.py | 335 +++++++----------- .../test_vertex_ai_live_passthrough.py | 326 ++++++++++++----- 2 files changed, 375 insertions(+), 286 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index e26f5f57532..e224f707b02 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -5,7 +5,10 @@ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough e Supports different modalities: text, audio, video, and web search. """ +from collections.abc import Mapping, Sequence from datetime import datetime +from itertools import chain +from types import MappingProxyType from typing import Any, Final from litellm._logging import verbose_proxy_logger @@ -15,8 +18,23 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( PassThroughEndpointLoggingTypedDict, ) -from litellm.types.utils import LlmProviders, ModelResponse, Usage -from litellm.utils import get_model_info +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + LlmProviders, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) + +_AGGREGATED_FIELDS: Final = frozenset( + { + "promptTokenCount", + "candidatesTokenCount", + "totalTokenCount", + "promptTokensDetails", + "candidatesTokensDetails", + } +) class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @@ -48,6 +66,56 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """Return the LLM provider name.""" return LlmProviders.VERTEX_AI + @staticmethod + def _resolve_detail_counts( + details: Sequence[Mapping[str, Any]], + declared_total: object, + ) -> tuple[tuple[str, int], ...]: + """ + Pair each of one turn's ``*TokensDetails`` entries with its token count. + + Live sometimes names the modality that carries the rest of a turn without a + ``tokenCount``, and reading the absent key as zero drops those tokens from the + breakdown, so real audio ends up priced as text. A lone unpriced entry therefore takes + whatever the turn's declared count leaves over. Two or more cannot be told apart, so + they are left out and the cost calculator charges the remainder as text. + """ + priced: Final = tuple( + (str(detail.get("modality", "TEXT")), count) + for detail in details + if isinstance(count := detail.get("tokenCount"), int) + ) + unpriced: Final = tuple( + str(detail.get("modality", "TEXT")) for detail in details if not isinstance(detail.get("tokenCount"), int) + ) + if len(unpriced) != 1 or not isinstance(declared_total, int): + return priced + residual: Final = declared_total - sum(count for _, count in priced) + return priced if residual <= 0 else (*priced, (unpriced[0], residual)) + + @staticmethod + def _sum_by_modality(counts: Sequence[tuple[str, int]]) -> Mapping[str, int]: + """Total the (modality, tokenCount) pairs of one or more turns per modality.""" + return MappingProxyType({modality: sum(c for m, c in counts if m == modality) for modality, _ in counts}) + + @staticmethod + def _merged_modality_totals( + snapshots: Sequence[Mapping[str, Any]], + count_key: str, + details_key: str, + ) -> Mapping[str, int]: + """Total every turn's per-modality counts, so the breakdown adds up the way the totals do.""" + return VertexAILivePassthroughLoggingHandler._sum_by_modality( + tuple( + chain.from_iterable( + VertexAILivePassthroughLoggingHandler._resolve_detail_counts( + snapshot.get(details_key) or [], snapshot.get(count_key) + ) + for snapshot in snapshots + ) + ) + ) + @staticmethod def _extract_usage_metadata_from_websocket_messages( websocket_messages: list[dict], @@ -55,175 +123,45 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Extract and aggregate usage metadata from a list of WebSocket messages. + Live emits one ``usageMetadata`` per turn and Google charges per turn for every token in + the session context window, which is the current turn's tokens plus all accumulated + tokens from previous turns, so the turns add up rather than restating each other. See + the Live API note under https://cloud.google.com/vertex-ai/generative-ai/pricing. + Args: websocket_messages: List of WebSocket messages from the Live API Returns: Dictionary containing aggregated usage metadata, or None if not found """ - all_usage_metadata: Final = [] + snapshots: Final = tuple( + message["usageMetadata"] + for message in websocket_messages + if isinstance(message, dict) and isinstance(message.get("usageMetadata"), dict) + ) - # Collect all usage metadata messages - for message in websocket_messages: - if isinstance(message, dict) and "usageMetadata" in message: - all_usage_metadata.append(message["usageMetadata"]) - - if not all_usage_metadata: + if not snapshots: return None - # If only one usage metadata, return it as-is - if len(all_usage_metadata) == 1: - return all_usage_metadata[0] - - # Aggregate multiple usage metadata messages - aggregated: Final[dict[str, Any]] = { - "promptTokenCount": 0, - "candidatesTokenCount": 0, - "totalTokenCount": 0, - "promptTokensDetails": [], - "candidatesTokensDetails": [], + prompt_totals: Final = VertexAILivePassthroughLoggingHandler._merged_modality_totals( + snapshots, "promptTokenCount", "promptTokensDetails" + ) + candidate_totals: Final = VertexAILivePassthroughLoggingHandler._merged_modality_totals( + snapshots, "candidatesTokenCount", "candidatesTokensDetails" + ) + return { + **{key: value for key, value in snapshots[0].items() if key not in _AGGREGATED_FIELDS}, + "promptTokenCount": sum(snapshot.get("promptTokenCount", 0) for snapshot in snapshots), + "candidatesTokenCount": sum(snapshot.get("candidatesTokenCount", 0) for snapshot in snapshots), + "totalTokenCount": sum(snapshot.get("totalTokenCount", 0) for snapshot in snapshots), + "promptTokensDetails": [ + {"modality": modality, "tokenCount": count} for modality, count in prompt_totals.items() if count > 0 + ], + "candidatesTokensDetails": [ + {"modality": modality, "tokenCount": count} for modality, count in candidate_totals.items() if count > 0 + ], } - # Aggregate token counts - for usage in all_usage_metadata: - aggregated["promptTokenCount"] += usage.get("promptTokenCount", 0) - aggregated["candidatesTokenCount"] += usage.get("candidatesTokenCount", 0) - aggregated["totalTokenCount"] += usage.get("totalTokenCount", 0) - - # Aggregate token details by modality - modality_totals: Final = {} - - for usage in all_usage_metadata: - # Process prompt tokens details - for detail in usage.get("promptTokensDetails", []): - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality not in modality_totals: - modality_totals[modality] = {"prompt": 0, "candidate": 0} - modality_totals[modality]["prompt"] += token_count - - # Process candidate tokens details - for detail in usage.get("candidatesTokensDetails", []): - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality not in modality_totals: - modality_totals[modality] = {"prompt": 0, "candidate": 0} - modality_totals[modality]["candidate"] += token_count - - # Convert aggregated modality totals back to details format - for modality, totals in modality_totals.items(): - if totals["prompt"] > 0: - aggregated["promptTokensDetails"].append({"modality": modality, "tokenCount": totals["prompt"]}) - if totals["candidate"] > 0: - aggregated["candidatesTokensDetails"].append({"modality": modality, "tokenCount": totals["candidate"]}) - - # Add any additional fields from the first usage metadata - first_usage: Final = all_usage_metadata[0] - for key, value in first_usage.items(): - if key not in aggregated: - aggregated[key] = value - - return aggregated - - @staticmethod - def _calculate_live_api_cost( - model: str, - usage_metadata: dict, - custom_llm_provider: str = "vertex_ai", - ) -> float: - """ - Calculate cost for Vertex AI Live API based on usage metadata. - - Args: - model: The model name (e.g., "gemini-2.0-flash-live-preview-04-09") - usage_metadata: Usage metadata from the Live API response - custom_llm_provider: The LLM provider (default: "vertex_ai") - - Returns: - Total cost in USD - """ - try: - # Get model pricing information - model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider) - - verbose_proxy_logger.debug("Vertex AI Live API model info for '%s': %s", model, model_info) - - # Check if pricing info is available - if not model_info or not model_info.get("input_cost_per_token"): - verbose_proxy_logger.error("No pricing info found for %s in local model pricing database", model) - return 0.0 - - total_cost = 0.0 - - # Extract token counts from usage metadata - prompt_token_count: Final = usage_metadata.get("promptTokenCount", 0) - candidates_token_count: Final = usage_metadata.get("candidatesTokenCount", 0) - - # Calculate base text token costs - input_cost_per_token: Final = model_info.get("input_cost_per_token", 0.0) - output_cost_per_token: Final = model_info.get("output_cost_per_token", 0.0) - - total_cost += prompt_token_count * input_cost_per_token - total_cost += candidates_token_count * output_cost_per_token - - # Handle modality-specific costs if present - prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", []) - candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", []) - - # Process prompt tokens by modality - for detail in prompt_tokens_details: - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality == "AUDIO": - audio_cost_per_token = model_info.get("input_cost_per_audio_token", 0.0) - total_cost += token_count * audio_cost_per_token - elif modality == "VIDEO": - # Video tokens are typically per second, but we'll treat as per token for now - video_cost_per_token = model_info.get("input_cost_per_video_per_second", 0.0) - total_cost += token_count * video_cost_per_token - # TEXT tokens are already handled above - - # Process candidate tokens by modality - for detail in candidates_tokens_details: - modality = detail.get("modality", "TEXT") - token_count = detail.get("tokenCount", 0) - - if modality == "AUDIO": - audio_cost_per_token = model_info.get("output_cost_per_audio_token", 0.0) - total_cost += token_count * audio_cost_per_token - elif modality == "VIDEO": - # Video tokens are typically per second, but we'll treat as per token for now - video_cost_per_token = model_info.get("output_cost_per_video_per_second", 0.0) - total_cost += token_count * video_cost_per_token - # TEXT tokens are already handled above - - # Handle web search costs if present - tool_use_prompt_token_count: Final = usage_metadata.get("toolUsePromptTokenCount", 0) - if tool_use_prompt_token_count > 0: - # Web search typically has a fixed cost per request - web_search_cost: Final = model_info.get("web_search_cost_per_request", 0.0) - if isinstance(web_search_cost, (int, float)) and web_search_cost > 0: - total_cost += web_search_cost - else: - # Fallback to token-based pricing for tool use - total_cost += tool_use_prompt_token_count * input_cost_per_token - - verbose_proxy_logger.debug( - f"Vertex AI Live API cost calculation - Model: {model}, " - f"Prompt tokens: {prompt_token_count}, " - f"Candidate tokens: {candidates_token_count}, " - f"Total cost: ${total_cost:.6f}" - ) - - return total_cost - - except Exception as e: - verbose_proxy_logger.error("Error calculating Vertex AI Live API cost: %s", e) - return 0.0 - @staticmethod def _create_usage_object_from_metadata( usage_metadata: dict, @@ -239,38 +177,37 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): Returns: LiteLLM Usage object """ - prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) - completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) - total_tokens: Final = usage_metadata.get("totalTokenCount", 0) + prompt_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( + VertexAILivePassthroughLoggingHandler._resolve_detail_counts( + usage_metadata.get("promptTokensDetails") or [], usage_metadata.get("promptTokenCount") + ) + ) + candidates_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( + VertexAILivePassthroughLoggingHandler._resolve_detail_counts( + usage_metadata.get("candidatesTokensDetails") or [], usage_metadata.get("candidatesTokenCount") + ) + ) - # Create modality-specific token details if available - prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", []) - candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", []) - - # Extract text tokens from details - text_prompt_tokens = 0 - text_completion_tokens = 0 - - for detail in prompt_tokens_details: - if detail.get("modality") == "TEXT": - text_prompt_tokens = detail.get("tokenCount", 0) - break - - for detail in candidates_tokens_details: - if detail.get("modality") == "TEXT": - text_completion_tokens = detail.get("tokenCount", 0) - break - - # If no text tokens found in details, use total counts - if text_prompt_tokens == 0: - text_prompt_tokens = prompt_tokens - if text_completion_tokens == 0: - text_completion_tokens = completion_tokens + prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) or sum(prompt_by_modality.values()) + completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) or sum(candidates_by_modality.values()) return Usage( - prompt_tokens=text_prompt_tokens, - completion_tokens=text_completion_tokens, - total_tokens=total_tokens, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=usage_metadata.get("totalTokenCount", 0) or (prompt_tokens + completion_tokens), + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=prompt_by_modality.get("TEXT"), + audio_tokens=prompt_by_modality.get("AUDIO"), + image_tokens=prompt_by_modality.get("IMAGE"), + video_tokens=prompt_by_modality.get("VIDEO"), + tool_use_tokens=usage_metadata.get("toolUsePromptTokenCount") or None, + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=candidates_by_modality.get("TEXT"), + audio_tokens=candidates_by_modality.get("AUDIO"), + image_tokens=candidates_by_modality.get("IMAGE"), + video_tokens=candidates_by_modality.get("VIDEO"), + ), ) def vertex_ai_live_passthrough_handler( @@ -316,13 +253,6 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): "kwargs": kwargs, } - # Calculate cost using Live API specific pricing - response_cost: Final = self._calculate_live_api_cost( - model=model, - usage_metadata=usage_metadata, - custom_llm_provider=custom_llm_provider, - ) - # Create Usage object for standard LiteLLM logging usage: Final = self._create_usage_object_from_metadata( usage_metadata=usage_metadata, @@ -339,8 +269,6 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): choices=[], ) - # Update kwargs with cost information - kwargs["response_cost"] = response_cost kwargs["model"] = model kwargs["custom_llm_provider"] = custom_llm_provider @@ -350,10 +278,13 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): allowed_pattern: Final = re.compile(r"^[A-Za-z0-9._\-:]+$") safe_model: Final = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]" verbose_proxy_logger.debug( - f"Vertex AI Live API passthrough cost tracking - " - f"Model: {safe_model}, Cost: ${response_cost:.6f}, " - f"Prompt tokens: {usage.prompt_tokens}, " - f"Completion tokens: {usage.completion_tokens}" + "Vertex AI Live API passthrough cost tracking - Model: %s, " + "Prompt tokens: %s %s, Completion tokens: %s %s", + safe_model, + usage.prompt_tokens, + usage.prompt_tokens_details, + usage.completion_tokens, + usage.completion_tokens_details, ) return { diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index e2eb6d0b68b..3b6a548b219 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -201,88 +201,247 @@ class TestVertexAILivePassthroughLoggingHandler: assert text_prompt["tokenCount"] == 10 assert audio_prompt["tokenCount"] == 10 - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" - ) - def test_calculate_cost_basic(self, mock_get_model_info, handler): - """Test basic cost calculation""" - mock_get_model_info.return_value = { - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000002, - } + def test_usage_carries_every_modality(self, handler): + """Regression: the Usage object reported only TEXT, so audio and image billed as nothing. + prompt_tokens must be the full count and the details must name each modality, + because the cost calculator prices audio and image from *_tokens_details. + """ usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, - } - - cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) - - # The cost calculation may include additional factors, so we check it's reasonable - expected_min_cost = (100 * 0.000001) + (50 * 0.000002) - assert cost >= expected_min_cost - assert cost > 0 - - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" - ) - def test_calculate_cost_with_audio(self, mock_get_model_info, handler): - """Test cost calculation with audio tokens""" - mock_get_model_info.return_value = { - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000002, - "input_cost_per_audio_token": 0.0001, - "output_cost_per_audio_token": 0.0002, - } - - usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, + "promptTokenCount": 1300, + "candidatesTokenCount": 124, + "totalTokenCount": 1424, "promptTokensDetails": [ - {"modality": "TEXT", "tokenCount": 80}, - {"modality": "AUDIO", "tokenCount": 20}, + {"modality": "TEXT", "tokenCount": 13}, + {"modality": "AUDIO", "tokenCount": 127}, + {"modality": "IMAGE", "tokenCount": 1160}, ], "candidatesTokensDetails": [ - {"modality": "TEXT", "tokenCount": 30}, - {"modality": "AUDIO", "tokenCount": 20}, + {"modality": "TEXT", "tokenCount": 29}, + {"modality": "AUDIO", "tokenCount": 95}, ], } - cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) + usage = handler._create_usage_object_from_metadata( + usage_metadata=usage_metadata, model="gemini-live-2.5-flash" + ) - # Should include both text and audio costs - assert cost > 0 - assert cost > (100 * 0.000001) + ( - 50 * 0.000002 - ) # Should be higher due to audio + assert usage.prompt_tokens == 1300, "the full prompt count must survive, not just its text share" + assert usage.completion_tokens == 124 + assert usage.prompt_tokens_details.text_tokens == 13 + assert usage.prompt_tokens_details.audio_tokens == 127 + assert usage.prompt_tokens_details.image_tokens == 1160 + assert usage.completion_tokens_details.text_tokens == 29 + assert usage.completion_tokens_details.audio_tokens == 95 - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" + def test_usage_sums_repeated_modality_entries(self, handler): + """A modality can appear more than once across aggregated turns; sum, don't overwrite.""" + usage = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": 40, + "candidatesTokenCount": 0, + "promptTokensDetails": [ + {"modality": "IMAGE", "tokenCount": 10}, + {"modality": "IMAGE", "tokenCount": 25}, + {"modality": "TEXT", "tokenCount": 5}, + ], + }, + model="gemini-live-2.5-flash", + ) + assert usage.prompt_tokens_details.image_tokens == 35 + assert usage.prompt_tokens_details.text_tokens == 5 + + NATIVE_AUDIO_MODEL = "gemini-live-2.5-flash-preview-native-audio-09-2025" + + # A four-turn native-audio session. Google charges per turn for the whole session context + # window, so the prompt side repeats the accumulated audio while the candidates side reports + # only that turn's own response. The last turn names AUDIO and omits its tokenCount, which is + # the shape Live really emits at the end of a spoken answer. + AUDIO_SESSION = ( + {"prompt": (14, 122), "candidates": (8, 20)}, + {"prompt": (21, 182), "candidates": (5, 50)}, + {"prompt": (24, 203), "candidates": (13, 27)}, + {"prompt": (24, 203), "candidates": (0, 3), "candidate_audio_token_count_missing": True}, ) - def test_calculate_cost_with_web_search(self, mock_get_model_info, handler): - """Test cost calculation with web search (tool use)""" - mock_get_model_info.return_value = { - "input_cost_per_token": 0.000001, - "output_cost_per_token": 0.000002, - "web_search_cost_per_request": 0.01, - } - usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, - "toolUsePromptTokenCount": 10, - } + @staticmethod + def _live_messages(turns): + """Wrap (text, audio) prompt/candidate pairs as the server messages a Live session emits.""" + return [{"type": "session.created", "session": {"id": "s"}}] + [ + { + "type": "response.done", + "usageMetadata": { + "promptTokenCount": sum(turn["prompt"]), + "candidatesTokenCount": sum(turn["candidates"]), + "totalTokenCount": sum(turn["prompt"]) + sum(turn["candidates"]), + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": turn["prompt"][0]}, + {"modality": "AUDIO", "tokenCount": turn["prompt"][1]}, + ], + "candidatesTokensDetails": ( + [{"modality": "AUDIO"}] + if turn.get("candidate_audio_token_count_missing") + else [ + {"modality": "TEXT", "tokenCount": turn["candidates"][0]}, + {"modality": "AUDIO", "tokenCount": turn["candidates"][1]}, + ] + ), + }, + } + for turn in turns + ] - cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata) + @staticmethod + def _session_usage(handler, mock_logging_obj, messages, model): + result = handler.vertex_ai_live_passthrough_handler( + websocket_messages=messages, + logging_obj=mock_logging_obj, + url_route="/vertex_ai/live", + start_time=datetime.now(), + end_time=datetime.now(), + request_body={}, + model=model, + ) + assert result["result"] is not None, "the handler must produce a usage-bearing response to bill" + return result["result"].usage - # Should include web search cost - expected_base_cost = (100 * 0.000001) + (50 * 0.000002) - # The web search cost might be handled differently, so just check it's reasonable - assert cost >= expected_base_cost - assert cost > 0 + @classmethod + def _session_cost(cls, handler, mock_logging_obj, messages, model): + from litellm.cost_calculator import completion_cost + from litellm.types.utils import ModelResponse + + usage = cls._session_usage(handler, mock_logging_obj, messages, model) + return completion_cost( + completion_response=ModelResponse( + id="x", object="chat.completion", created=0, model=model, usage=usage, choices=[] + ), + model=f"vertex_ai/{model}", + custom_llm_provider="vertex_ai", + call_type="acompletion", + ) + + @classmethod + def _expected_session_cost(cls, turns): + from litellm.utils import get_model_info + + info = get_model_info(model=cls.NATIVE_AUDIO_MODEL, custom_llm_provider="vertex_ai") + return ( + sum(turn["prompt"][0] for turn in turns) * info["input_cost_per_token"] + + sum(turn["prompt"][1] for turn in turns) * info["input_cost_per_audio_token"] + + sum(turn["candidates"][0] for turn in turns) * info["output_cost_per_token"] + + sum(turn["candidates"][1] for turn in turns) * info["output_cost_per_audio_token"] + ) + + def test_every_turn_of_a_session_is_billed(self, handler, mock_logging_obj): + """Google charges per turn for the whole context window, so every turn adds to the bill. + + Billing one snapshot instead gives away all the other turns: on this session the + largest single turn is well under the session total, and its share of the audio is + priced 6x the text rate, so the gap is money rather than rounding. + """ + turns = self.AUDIO_SESSION[:3] + cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + + assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) + widest_single_turn = max(self._expected_session_cost([turn]) for turn in turns) + assert cost > widest_single_turn, "billing one snapshot drops every other turn of the session" + + def test_audio_named_without_a_token_count_bills_at_the_audio_rate(self, handler, mock_logging_obj): + """Live can name the modality carrying the rest of a turn and omit its tokenCount. + + Reading the absent key as zero left those tokens inside candidatesTokenCount but outside + the breakdown, so the calculator charged real speech at the text output rate. At this + entry's rates the last turn's 3 audio tokens are $0.0000360 rather than $0.0000060. + """ + turns = self.AUDIO_SESSION + usage = self._session_usage(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + + assert usage.completion_tokens_details.audio_tokens == 100, "the unpriced entry takes the turn's residual" + assert usage.completion_tokens_details.text_tokens == 26 + assert usage.completion_tokens == 126 + + cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) + + def test_server_side_tool_use_prompt_tokens_are_reported(self, handler, mock_logging_obj): + """toolUsePromptTokenCount was dropped, so a grounded session logged fewer tokens than it used. + + It is reported, not billed. Nothing in the shared Gemini input-cost path prices + tool-use tokens, and folding them into prompt_tokens here would suppress that + path's cache-overlap correction and raise the bill instead. + """ + messages = self._live_messages(self.AUDIO_SESSION[:1]) + grounded = [dict(message) for message in messages] + grounded[-1]["usageMetadata"] = {**grounded[-1]["usageMetadata"], "toolUsePromptTokenCount": 500} + + usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) + assert usage.prompt_tokens_details.tool_use_tokens == 500 + + plain_cost = self._session_cost(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) + grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) + assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill" + + @pytest.mark.parametrize( + "label,prompt_details,candidate_details", + [ + ("text only", [("TEXT", 6)], [("TEXT", 2)]), + ("audio in", [("TEXT", 13), ("AUDIO", 127)], [("TEXT", 18)]), + ("image in", [("TEXT", 10), ("IMAGE", 258)], [("TEXT", 24)]), + ("frames in", [("TEXT", 11), ("IMAGE", 1032)], [("TEXT", 26)]), + ("audio both ways", [("TEXT", 13), ("AUDIO", 127)], [("TEXT", 29), ("AUDIO", 95)]), + ], + ) + def test_live_session_bills_each_modality_at_its_own_rate(self, handler, label, prompt_details, candidate_details): + """Every payload here is a real Vertex Live session's usageMetadata. + + Before the fix these billed the text share only, from 1x (text) to 55x under. + The expected amount is derived from the entry's own rates rather than hardcoded, + so this stays correct as prices move, and it is asserted exactly, so dropping a + modality and double-charging one both fail. + """ + from litellm.cost_calculator import completion_cost + from litellm.types.utils import ModelResponse + from litellm.utils import get_model_info + + model = self.NATIVE_AUDIO_MODEL + info = get_model_info(model=model, custom_llm_provider="vertex_ai") + + text_in = info["input_cost_per_token"] + audio_in = info.get("input_cost_per_audio_token") or text_in + image_in = info.get("input_cost_per_image_token") or text_in + text_out = info["output_cost_per_token"] + audio_out = info.get("output_cost_per_audio_token") or text_out + rate_in = {"TEXT": text_in, "AUDIO": audio_in, "IMAGE": image_in} + rate_out = {"TEXT": text_out, "AUDIO": audio_out} + + expected = sum(c * rate_in[m] for m, c in prompt_details) + sum(c * rate_out[m] for m, c in candidate_details) + + usage = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": sum(c for _, c in prompt_details), + "candidatesTokenCount": sum(c for _, c in candidate_details), + "promptTokensDetails": [{"modality": m, "tokenCount": c} for m, c in prompt_details], + "candidatesTokensDetails": [{"modality": m, "tokenCount": c} for m, c in candidate_details], + }, + model=model, + ) + + cost = completion_cost( + completion_response=ModelResponse( + id="x", object="chat.completion", created=0, model=model, usage=usage, choices=[] + ), + model=f"vertex_ai/{model}", + custom_llm_provider="vertex_ai", + call_type="acompletion", + ) + + assert cost == pytest.approx(expected, rel=1e-9), label + + text_only = sum(c for m, c in prompt_details if m == "TEXT") * text_in + sum( + c for m, c in candidate_details if m == "TEXT" + ) * text_out + if any(m != "TEXT" for m, _ in prompt_details + candidate_details) and audio_in != text_in: + assert cost > text_only, f"{label}: non-text modalities must add cost" def test_vertex_ai_live_passthrough_handler_integration( self, handler, mock_logging_obj, sample_websocket_messages @@ -540,25 +699,24 @@ class TestVertexAILivePassthroughErrorHandling: result = handler._extract_usage_metadata_from_websocket_messages(messages) assert result is None - @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info" - ) - def test_cost_calculation_with_missing_model_info(self, mock_get_model_info): - """Test cost calculation when model info is missing""" + def test_usage_without_modality_details(self): + """Older payloads carry only the totals; fall back to them rather than reporting zero.""" handler = VertexAILivePassthroughLoggingHandler() - # Mock missing model info - mock_get_model_info.return_value = {} + usage = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "totalTokenCount": 150, + }, + model="unknown-model", + ) - usage_metadata = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, - } - - # Should not raise an exception, should return 0 or handle gracefully - cost = handler._calculate_live_api_cost("unknown-model", usage_metadata) - assert cost == 0.0 + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 50 + assert usage.total_tokens == 150 + assert usage.prompt_tokens_details.audio_tokens is None + assert usage.prompt_tokens_details.image_tokens is None def test_handler_with_none_websocket_messages(self, mock_logging_obj): """Test handler with None websocket messages""" From da73896ec38d7cc5606d1563d87ddd7db31a4d2f Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 04:43:50 -0400 Subject: [PATCH 011/107] refactor(vertex-live): type the usage helpers without Any The two helpers this branch adds took Sequence[Mapping[str, Any]], which the repo forbids, and only typechecked because Any is compatible with everything. Both now take Mapping[str, object] and the raw *TokensDetails value is narrowed to its mapping entries at each of the three call sites. TypedDicts are the wrong tool here: _merged_modality_totals reads count_key and details_key as runtime strings, and the aggregation deliberately passes unknown keys straight through, so both need a mapping whose keys are not literals. The narrowing is not cosmetic. The handler's only failure path returns no result at all, so a *TokensDetails value that was not a list of objects used to raise while being read and cost the whole session its bill. --- ...tex_ai_live_passthrough_logging_handler.py | 18 +++++++---- .../test_vertex_ai_live_passthrough.py | 30 +++++++++++++++++++ 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index e224f707b02..72e8c3f7657 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -9,7 +9,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime from itertools import chain from types import MappingProxyType -from typing import Any, Final +from typing import Final from litellm._logging import verbose_proxy_logger from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( @@ -37,6 +37,11 @@ _AGGREGATED_FIELDS: Final = frozenset( ) +def _detail_entries(raw: object) -> tuple[Mapping[str, object], ...]: + """Narrow one turn's ``*TokensDetails`` value to the entries that are actually shaped like one.""" + return tuple(entry for entry in raw if isinstance(entry, Mapping)) if isinstance(raw, Sequence) else () + + class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough. @@ -68,7 +73,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @staticmethod def _resolve_detail_counts( - details: Sequence[Mapping[str, Any]], + details: Sequence[Mapping[str, object]], declared_total: object, ) -> tuple[tuple[str, int], ...]: """ @@ -100,7 +105,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @staticmethod def _merged_modality_totals( - snapshots: Sequence[Mapping[str, Any]], + snapshots: Sequence[Mapping[str, object]], count_key: str, details_key: str, ) -> Mapping[str, int]: @@ -109,7 +114,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): tuple( chain.from_iterable( VertexAILivePassthroughLoggingHandler._resolve_detail_counts( - snapshot.get(details_key) or [], snapshot.get(count_key) + _detail_entries(snapshot.get(details_key)), snapshot.get(count_key) ) for snapshot in snapshots ) @@ -179,12 +184,13 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ prompt_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( VertexAILivePassthroughLoggingHandler._resolve_detail_counts( - usage_metadata.get("promptTokensDetails") or [], usage_metadata.get("promptTokenCount") + _detail_entries(usage_metadata.get("promptTokensDetails")), usage_metadata.get("promptTokenCount") ) ) candidates_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality( VertexAILivePassthroughLoggingHandler._resolve_detail_counts( - usage_metadata.get("candidatesTokensDetails") or [], usage_metadata.get("candidatesTokenCount") + _detail_entries(usage_metadata.get("candidatesTokensDetails")), + usage_metadata.get("candidatesTokenCount"), ) ) diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index 3b6a548b219..ba9167e2b13 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -381,6 +381,36 @@ class TestVertexAILivePassthroughLoggingHandler: grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill" + def test_a_malformed_details_entry_does_not_cost_the_whole_session(self, handler, mock_logging_obj): + """A ``*TokensDetails`` value that is not a list of objects must not take the session down. + + The handler's only error path returns no result at all, so one odd frame used to throw + while reading it and the whole session billed nothing. The good turns still bill. + """ + turns = self.AUDIO_SESSION[:3] + messages = self._live_messages(turns) + mangled = [dict(message) for message in messages] + mangled[1]["usageMetadata"] = {**mangled[1]["usageMetadata"], "promptTokensDetails": "TEXT"} + + usage = self._session_usage(handler, mock_logging_obj, mangled, self.NATIVE_AUDIO_MODEL) + + surviving = turns[1:] + assert usage.prompt_tokens_details.audio_tokens == sum(turn["prompt"][1] for turn in surviving) + assert usage.prompt_tokens_details.text_tokens == sum(turn["prompt"][0] for turn in surviving) + assert usage.prompt_tokens == sum(sum(turn["prompt"]) for turn in turns), "the totals still cover every turn" + + direct = handler._create_usage_object_from_metadata( + usage_metadata={ + "promptTokenCount": 40, + "candidatesTokenCount": 12, + "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 40}, "AUDIO"], + "candidatesTokensDetails": {"modality": "TEXT", "tokenCount": 12}, + }, + model=self.NATIVE_AUDIO_MODEL, + ) + assert direct.prompt_tokens_details.audio_tokens == 40, "the well-formed entry beside a bad one still counts" + assert direct.completion_tokens == 12 + @pytest.mark.parametrize( "label,prompt_details,candidate_details", [ From b051aa713a3d4c044fb9453935879f09d91eea69 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 04:54:13 -0400 Subject: [PATCH 012/107] fix(vertex-live): sum tool-use prompt tokens across a session's turns toolUsePromptTokenCount was the one prompt-side total not named in _AGGREGATED_FIELDS, so it rode the unknown-key pass-through and took the first frame's value while promptTokenCount, candidatesTokenCount and totalTokenCount beside it were summed. Live's frames grow over a session, so the first frame is the smallest number in the series and a grounded session under-reported its tool-use tokens by everything after turn one. It is now summed like its three neighbours. This is reporting only, and pricing these tokens is deliberately left out. Google charges tool-use prompt tokens at the input token rate, but generic_cost_per_token reads the input bill out of prompt_tokens_details and only falls back to prompt_tokens when the details carry no text or a cache hit overlaps them. Measured on the native-audio entry with 500 tool-use tokens: adding them to prompt_tokens moves an ordinary Live turn's bill by $0.0000000000, and on a turn with a cache hit it moves it by $0.0002650000 where the tokens are worth $0.0002500000, because it perturbs the cache-overlap correction. Pricing them belongs beside the modality terms in the shared input-cost path, in its own change that fixes the same latent no-op on the ordinary Gemini path. Not verified against a live capture: no Vertex Live session we have captured reported toolUsePromptTokenCount at all, so the summing convention is inferred from the three prompt-side totals that accumulate the same way. --- ...tex_ai_live_passthrough_logging_handler.py | 2 + .../test_vertex_ai_live_passthrough.py | 46 ++++++++++++++----- 2 files changed, 37 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 72e8c3f7657..2cd336db3d6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -31,6 +31,7 @@ _AGGREGATED_FIELDS: Final = frozenset( "promptTokenCount", "candidatesTokenCount", "totalTokenCount", + "toolUsePromptTokenCount", "promptTokensDetails", "candidatesTokensDetails", } @@ -159,6 +160,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): "promptTokenCount": sum(snapshot.get("promptTokenCount", 0) for snapshot in snapshots), "candidatesTokenCount": sum(snapshot.get("candidatesTokenCount", 0) for snapshot in snapshots), "totalTokenCount": sum(snapshot.get("totalTokenCount", 0) for snapshot in snapshots), + "toolUsePromptTokenCount": sum(snapshot.get("toolUsePromptTokenCount", 0) for snapshot in snapshots), "promptTokensDetails": [ {"modality": modality, "tokenCount": count} for modality, count in prompt_totals.items() if count > 0 ], diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index ba9167e2b13..bfff52de673 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -363,22 +363,46 @@ class TestVertexAILivePassthroughLoggingHandler: cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) - def test_server_side_tool_use_prompt_tokens_are_reported(self, handler, mock_logging_obj): - """toolUsePromptTokenCount was dropped, so a grounded session logged fewer tokens than it used. + TOOL_USE_PER_TURN = (100, 250, 400) - It is reported, not billed. Nothing in the shared Gemini input-cost path prices - tool-use tokens, and folding them into prompt_tokens here would suppress that - path's cache-overlap correction and raise the bill instead. + def _grounded_messages(self): + """The three-turn session again, with each turn's own toolUsePromptTokenCount attached.""" + messages = self._live_messages(self.AUDIO_SESSION[:3]) + head, turns = messages[0], messages[1:] + return [head] + [ + {**message, "usageMetadata": {**message["usageMetadata"], "toolUsePromptTokenCount": tool_use}} + for message, tool_use in zip(turns, self.TOOL_USE_PER_TURN) + ] + + def test_server_side_tool_use_prompt_tokens_are_summed_over_the_session(self, handler, mock_logging_obj): + """toolUsePromptTokenCount rode the unknown-key pass-through, so it took the first turn only. + + Every other total beside it is summed across the session, and the first turn is the + smallest number in the series, so a grounded session logged far fewer tool-use tokens + than it used. This session's turns are deliberately distinct, so 750 can only come from + summing: first-turn selection gives 100, last-turn or max gives 400. """ - messages = self._live_messages(self.AUDIO_SESSION[:1]) - grounded = [dict(message) for message in messages] - grounded[-1]["usageMetadata"] = {**grounded[-1]["usageMetadata"], "toolUsePromptTokenCount": 500} + grounded = self._grounded_messages() usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) - assert usage.prompt_tokens_details.tool_use_tokens == 500 + assert usage.prompt_tokens_details.tool_use_tokens == sum(self.TOOL_USE_PER_TURN) - plain_cost = self._session_cost(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) - grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): + """Deliberate boundary: these tokens are reported here, and priced nowhere. + + generic_cost_per_token reads the input bill out of prompt_tokens_details, and falls + back to prompt_tokens only when the details carry no text or a cache hit overlaps them, + so adding tool-use tokens to prompt_tokens is worth nothing on an ordinary Live turn and + over-charges against the cache-overlap correction when it is not. Pricing them belongs + in the shared input-cost path, beside the modality terms that already read the details. + """ + turns = self.AUDIO_SESSION[:3] + plain_cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + grounded_cost = self._session_cost( + handler, mock_logging_obj, self._grounded_messages(), self.NATIVE_AUDIO_MODEL + ) + + assert plain_cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill" def test_a_malformed_details_entry_does_not_cost_the_whole_session(self, handler, mock_logging_obj): From 22e7c7a5338a96009607c612804fe90c0dad9a1a Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 06:25:07 -0400 Subject: [PATCH 013/107] test(vertex-live): type the session helpers in the Live passthrough tests The four session helpers this PR added were unannotated. Typing them needs a name for the (text, audio) pair each turn carries, so _LiveTurn is a TypedDict rather than a Mapping union that would leave sum() over a prompt pair ill-typed, and AUDIO_SESSION is declared with it. The message list reuses list[dict[str, object]], the annotation the passthrough already uses where it collects those messages --- .../test_vertex_ai_live_passthrough.py | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index bfff52de673..b4d0a6c06e5 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -6,12 +6,14 @@ including the logging handler, cost tracking, and WebSocket message processing. """ import json +from collections.abc import Sequence from datetime import datetime from unittest.mock import AsyncMock, Mock, patch, MagicMock from typing import Dict, List, Any, Optional import pytest import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict # Add the parent directory to the system path @@ -22,10 +24,16 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.types.utils import LlmProviders +from litellm.types.utils import LlmProviders, Usage from litellm.proxy._types import UserAPIKeyAuth +class _LiveTurn(TypedDict): + prompt: ReadOnly[tuple[int, int]] + candidates: ReadOnly[tuple[int, int]] + candidate_audio_token_count_missing: NotRequired[ReadOnly[bool]] + + class TestVertexAILivePassthroughLoggingHandler: """Test the Vertex AI Live Passthrough Logging Handler""" @@ -257,7 +265,7 @@ class TestVertexAILivePassthroughLoggingHandler: # window, so the prompt side repeats the accumulated audio while the candidates side reports # only that turn's own response. The last turn names AUDIO and omits its tokenCount, which is # the shape Live really emits at the end of a spoken answer. - AUDIO_SESSION = ( + AUDIO_SESSION: tuple[_LiveTurn, ...] = ( {"prompt": (14, 122), "candidates": (8, 20)}, {"prompt": (21, 182), "candidates": (5, 50)}, {"prompt": (24, 203), "candidates": (13, 27)}, @@ -265,7 +273,7 @@ class TestVertexAILivePassthroughLoggingHandler: ) @staticmethod - def _live_messages(turns): + def _live_messages(turns: Sequence[_LiveTurn]) -> list[dict[str, object]]: """Wrap (text, audio) prompt/candidate pairs as the server messages a Live session emits.""" return [{"type": "session.created", "session": {"id": "s"}}] + [ { @@ -292,7 +300,12 @@ class TestVertexAILivePassthroughLoggingHandler: ] @staticmethod - def _session_usage(handler, mock_logging_obj, messages, model): + def _session_usage( + handler: VertexAILivePassthroughLoggingHandler, + mock_logging_obj: MagicMock, + messages: list[dict[str, object]], + model: str, + ) -> Usage: result = handler.vertex_ai_live_passthrough_handler( websocket_messages=messages, logging_obj=mock_logging_obj, @@ -306,7 +319,13 @@ class TestVertexAILivePassthroughLoggingHandler: return result["result"].usage @classmethod - def _session_cost(cls, handler, mock_logging_obj, messages, model): + def _session_cost( + cls, + handler: VertexAILivePassthroughLoggingHandler, + mock_logging_obj: MagicMock, + messages: list[dict[str, object]], + model: str, + ) -> float: from litellm.cost_calculator import completion_cost from litellm.types.utils import ModelResponse @@ -321,7 +340,7 @@ class TestVertexAILivePassthroughLoggingHandler: ) @classmethod - def _expected_session_cost(cls, turns): + def _expected_session_cost(cls, turns: Sequence[_LiveTurn]) -> float: from litellm.utils import get_model_info info = get_model_info(model=cls.NATIVE_AUDIO_MODEL, custom_llm_provider="vertex_ai") From b1262b05f470f550ea0a2d879449f4fd966c4a1d Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 22:05:02 -0400 Subject: [PATCH 014/107] fix(gemini-live): count grounding requests so Live sessions carry their query fee Live reports grounding in the server frames and never in usageMetadata, so nothing set the counter the cost path reads and the per-query charge was missing from every grounded session. Google bills a grounded Live prompt on top of its tokens, and that fee dwarfs the token cost, so a non-zero spend check could never catch it. Both Live surfaces now read serverContent.groundingMetadata where they build usage, and reuse the chat path's own classifier so web search and Maps keep their separate SKUs rather than being counted together. Separately, a client sending turn_detection: null reached a membership test against None and took the session down with no traceback, while the branch immediately above already guards for it. Live emits grounding and usage on the same frame, verified against Vertex directly, so the realtime counter is set where usage is built. (cherry picked from commit c997436be34beb2e84f8468286b8016c195eca92) (cherry picked from commit 26c8d4822fc1c5c44fe8f72f4f57473a2ce1acbf) --- .../llms/gemini/realtime/transformation.py | 20 +++- ...tex_ai_live_passthrough_logging_handler.py | 32 ++++++- .../test_vertex_ai_live_passthrough.py | 68 +++++++++++++ .../test_gemini_realtime_transformation.py | 95 ++++++++++++++++++- 4 files changed, 212 insertions(+), 3 deletions(-) diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index c92af7de145..79985569c5f 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -115,6 +115,19 @@ def _parse_setup(session_configuration_request: str) -> BidiGenerateContentSetup return envelope.get("setup", empty_setup) +def _grounding_metadata_from_frame(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + """Read ``serverContent.groundingMetadata`` off the frame that carries the turn's usage. + + Live reports grounding in the server frames rather than in ``usageMetadata``, and it emits both + on the same frame, so the per-query charge is countable at the point usage is built. + """ + server_content: Final = frame.get("serverContent") + if not isinstance(server_content, Mapping): + return () + metadata: Final = server_content.get("groundingMetadata") + return (metadata,) if isinstance(metadata, Mapping) else () + + # Google bills Live transcription at an estimated 25 audio tokens/sec of input and # 175 text tokens/min of output (ai.google.dev/gemini-api/docs/pricing). GEMINI_LIVE_TRANSCRIBE_AUDIO_TOKENS_PER_SECOND: Final = 25 @@ -323,7 +336,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): ) elif key == "input_audio_transcription" and value is not None: optional_params["inputAudioTranscription"] = {} - elif key == "turn_detection": + elif key == "turn_detection" and value is not None: value_typed = cast(OpenAIRealtimeTurnDetection, value) if ( isinstance(value_typed, dict) @@ -1049,6 +1062,11 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): {**cast(dict, message), "usageMetadata": resolved_usage_metadata}, ), ) + grounding_metadata: Final = _grounding_metadata_from_frame(message) + if grounding_metadata: + VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet + _chat_completion_usage, grounding_metadata + ) else: _chat_completion_usage = get_empty_usage() diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 2cd336db3d6..0e2eb60704d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -43,6 +43,23 @@ def _detail_entries(raw: object) -> tuple[Mapping[str, object], ...]: return tuple(entry for entry in raw if isinstance(entry, Mapping)) if isinstance(raw, Sequence) else () +def _grounding_metadata(websocket_messages: Sequence[object]) -> tuple[Mapping[str, object], ...]: + """Collect every ``serverContent.groundingMetadata`` a session emitted. + + Live reports grounding in the server frames, never in ``usageMetadata``, so the per-query + charge has to be counted here rather than derived from the token totals. + """ + return tuple( + metadata + for message in websocket_messages + if isinstance(message, Mapping) + for server_content in (message.get("serverContent"),) + if isinstance(server_content, Mapping) + for metadata in (server_content.get("groundingMetadata"),) + if isinstance(metadata, Mapping) + ) + + class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough. @@ -173,6 +190,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): def _create_usage_object_from_metadata( usage_metadata: dict, model: str, + grounding_metadata: Sequence[Mapping[str, object]] = (), ) -> Usage: """ Create a LiteLLM Usage object from Live API usage metadata. @@ -180,6 +198,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): Args: usage_metadata: Usage metadata from the Live API response model: The model name + grounding_metadata: Every ``serverContent.groundingMetadata`` the session emitted, so + Search and Maps grounding carry their per-query charge Returns: LiteLLM Usage object @@ -199,7 +219,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) or sum(prompt_by_modality.values()) completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) or sum(candidates_by_modality.values()) - return Usage( + usage: Final = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=usage_metadata.get("totalTokenCount", 0) or (prompt_tokens + completion_tokens), @@ -217,6 +237,15 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): video_tokens=candidates_by_modality.get("VIDEO"), ), ) + if grounding_metadata: + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet + usage, grounding_metadata + ) + return usage def vertex_ai_live_passthrough_handler( self, @@ -264,6 +293,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): # Create Usage object for standard LiteLLM logging usage: Final = self._create_usage_object_from_metadata( usage_metadata=usage_metadata, + grounding_metadata=_grounding_metadata(websocket_messages), model=model, ) diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index b4d0a6c06e5..1815ff134aa 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -406,6 +406,74 @@ class TestVertexAILivePassthroughLoggingHandler: usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) assert usage.prompt_tokens_details.tool_use_tokens == sum(self.TOOL_USE_PER_TURN) + @staticmethod + def _grounding_frame(metadata: dict[str, object]) -> dict[str, object]: + """One server frame carrying grounding metadata, the way Live reports it.""" + return {"type": "response.done", "serverContent": {"groundingMetadata": metadata}} + + def test_web_grounding_is_counted_so_it_can_be_billed(self, handler, mock_logging_obj): + """Live reports grounding in the server frames and never in usageMetadata. + + Nothing read those frames, so web_search_requests stayed unset and the cost path's only + trigger for the per-query grounding charge never fired. Google bills a grounded Live + prompt on top of its tokens, so the whole fee was missing from the bill. + """ + messages = [ + self._grounding_frame( + { + "webSearchQueries": ["who won the 2026 world cup final"], + "groundingChunks": [{"web": {"uri": "https://example.com"}}], + } + ), + *self._live_messages(self.AUDIO_SESSION[:1]), + ] + + usage = self._session_usage(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) + + assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query" + assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None + + def test_maps_grounding_is_counted_under_its_own_sku(self, handler, mock_logging_obj): + """Maps grounding is a separate SKU from web search, so it needs its own counter. + + A maps-only turn carries grounding chunks but no webSearchQueries, so counting queries + alone would report nothing and bill nothing. + """ + messages = [ + self._grounding_frame({"groundingChunks": [{"maps": {"placeId": "abc123"}}]}), + *self._live_messages(self.AUDIO_SESSION[:1]), + ] + + usage = self._session_usage(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) + + assert usage.prompt_tokens_details.google_maps_grounding_requests == 1 + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + + def test_an_ungrounded_session_reports_no_grounding(self, handler, mock_logging_obj): + """The counters must stay absent when no tool ran, or every session pays a grounding fee.""" + usage = self._session_usage( + handler, mock_logging_obj, self._live_messages(self.AUDIO_SESSION[:1]), self.NATIVE_AUDIO_MODEL + ) + + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None + + def test_grounding_adds_its_query_fee_to_the_session_bill(self, handler, mock_logging_obj): + """The counter only matters if it reaches the bill, so assert against the cost, not the field. + + Same tokens either way: the difference between the two sessions is the grounding fee alone. + """ + turns = self.AUDIO_SESSION[:1] + plain = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + grounded = self._session_cost( + handler, + mock_logging_obj, + [self._grounding_frame({"webSearchQueries": ["q"]}), *self._live_messages(turns)], + self.NATIVE_AUDIO_MODEL, + ) + + assert grounded > plain, "a grounded session must cost more than the same tokens ungrounded" + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): """Deliberate boundary: these tokens are reported here, and priced nowhere. diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index 2b3b6343fad..ee984cc7e1f 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,11 +1,16 @@ import json -from unittest.mock import MagicMock +from collections.abc import Mapping +from typing import cast +from unittest.mock import AsyncMock, MagicMock, patch import pytest import litellm from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig +from litellm.types.llms.gemini import BidiGenerateContentServerMessage +from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents +from litellm.types.utils import Usage def test_gemini_realtime_transformation_session_created(): @@ -2178,3 +2183,91 @@ def test_unbilled_usage_on_session_close_flushes_trailing_audio(patch_gemini_tra } assert usage == expected assert config.unbilled_usage_on_session_close("gemini-3.5-transcribe-live") is None + + +def _grounded_live_frame(grounding_metadata: Mapping[str, object] | None) -> Mapping[str, object]: + """One Live server frame. Grounding metadata and usageMetadata arrive together, as Vertex sends them.""" + from typing import Final + + server_content: Final = { + "turnComplete": True, + **({} if grounding_metadata is None else {"groundingMetadata": grounding_metadata}), + } + return { + "serverContent": server_content, + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 157, + "totalTokenCount": 176, + "promptTokensDetails": ({"modality": "TEXT", "tokenCount": 19},), + "candidatesTokensDetails": ({"modality": "AUDIO", "tokenCount": 157},), + }, + } + + +def _usage_built_for_response_done(message: Mapping[str, object]) -> Usage: + """Capture the chat-completion Usage transform_response_done_event builds, before it is bridged. + + The Usage object is local to the method, so the bridge call is the only place it is observable. + """ + from typing import Final + + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + captured: Final[list[Usage]] = [] # mutable-ok: a spy has to accumulate what it observes + original: Final = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage + + def _spy(usage: Usage) -> object: + captured.append(usage) + return original(usage) + + config: Final = GeminiRealtimeConfig() + with patch.object( + LiteLLMCompletionResponsesConfig, + "_transform_chat_completion_usage_to_responses_usage", + staticmethod(_spy), + ): + config.transform_response_done_event( + message=cast( # cast-ok: a test fixture stands in for the server frame TypedDict + BidiGenerateContentServerMessage, message + ), + current_response_id="resp_grounding", + current_conversation_id="conv_grounding", + output_items=None, + ) + assert captured, "response.done must build a Usage object" + return captured[0] + + +def test_gemini_realtime_response_done_counts_web_grounding(): + """Regression: Live reports grounding in the server frames and never in usageMetadata. + + Nothing read those frames on the realtime path, so web_search_requests stayed unset and the + cost path's only trigger for Google's per-query grounding charge never fired. + + Scope boundary, deliberate: this asserts the counter on the Usage object that response.done is + built from, not on the emitted event. The Responses usage bridge copies a fixed allow-list of + detail fields and drops the rest, so the counter does not reach response.done yet. Widening + that bridge is a separate change; do not read this test as proving end-to-end billing. + """ + usage = _usage_built_for_response_done( + _grounded_live_frame( + { + "webSearchQueries": ["who won the 2026 world cup final"], + "groundingChunks": [{"web": {"uri": "https://example.com"}}], + } + ) + ) + + assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query" + assert usage.prompt_tokens_details.text_tokens == 19, "the modality breakdown must survive alongside it" + + +def test_gemini_realtime_response_done_reports_no_grounding_when_none_ran(): + """The counter must stay unset on an ordinary turn, or every session pays a grounding fee.""" + usage = _usage_built_for_response_done(_grounded_live_frame(None)) + + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None From 9580b89bb16d9946152aa23a7a4d7d139700d403 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 23:26:31 -0400 Subject: [PATCH 015/107] fix(vertex-live): resolve a Live setup model before logging reads it A client that named a bare gateway alias logged the session as "unknown" and billed nothing, because the model was read off the raw setup frame and the extractor only yields a name when the string already contains "/models/". The rewriter qualifies that same model a few lines later for the upstream, so the supported client form, an alias, was the one that went unbilled. Resolving through the rewriter first means the real model reaches the logging object, and from there the cost map. A route with no rewriter, which is every non-Live passthrough, hands the frame over untouched. (cherry picked from commit 573982803df612fd94144e2e06dd647f8530d4e8) --- .../pass_through_endpoints.py | 23 ++++- .../test_pass_through_endpoints.py | 94 +++++++++++++++++++ 2 files changed, 116 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index ea4ede7e513..5a2f9c391a2 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2090,6 +2090,22 @@ def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Calla return json.dumps({**message, "setup": {**setup, "model": rewritten_model}}) # mutable-ok: one-shot json payload +def _resolved_vertex_live_setup( + setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None +) -> Mapping[str, object]: + """ + Give the model extractor the same fully qualified path the upstream will receive. + + Clients may name a bare gateway alias, which the rewriter turns into a ``projects/...`` path before + it reaches Vertex. The extractor only reads a path containing ``/models/``, so running it on the raw + frame logs the session as ``unknown`` at no cost, which is precisely the supported client form + """ + setup_model: Final = setup_data.get("model") + if setup_model_rewriter is None or not isinstance(setup_model, str): + return setup_data + return {**setup_data, "model": setup_model_rewriter(setup_model)} + + def _truncated_close_reason(reason: str) -> str: """ Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character @@ -2314,7 +2330,12 @@ async def websocket_passthrough_request( setup_data, ) if isinstance(setup_data, dict) and "model" in setup_data: - extracted_model = _extract_model_from_vertex_ai_setup(setup_data) + # Resolve the alias first: a client may name a bare gateway model, + # which carries no "/models/" for the extractor to read, so reading + # the raw frame leaves the session logged as "unknown" and unbilled. + extracted_model = _extract_model_from_vertex_ai_setup( + _resolved_vertex_live_setup(setup_data, setup_model_rewriter) + ) if extracted_model: kwargs["model"] = extracted_model kwargs["custom_llm_provider"] = "vertex_ai-language-models" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d57bed430c1..11066d4ed38 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -5030,6 +5030,100 @@ async def test_websocket_passthrough_rewrites_gateway_alias_setup_model(): assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash" +@pytest.mark.parametrize( + "setup_model", + ["gemini-live-2.5-flash", "models/gemini-live-2.5-flash", "publishers/google/models/gemini-live-2.5-flash"], +) +def test_vertex_live_setup_model_resolves_before_extraction(setup_model): + """A bare gateway alias left the session logged as ``unknown`` at zero cost. + + The model was read off the raw client frame, and the extractor only yields a name when the string + already contains ``/models/``. The rewriter qualifies it a few lines later for the upstream, so a + client that addressed the gateway the documented way, by alias, logged no model and therefore + resolved no cost-map entry. Resolving first is what puts the real name on the logging object. + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _build_vertex_live_setup_model_rewriter, + ) + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _extract_model_from_vertex_ai_setup, + _resolved_vertex_live_setup, + ) + + rewriter = _build_vertex_live_setup_model_rewriter( + vertex_project="proj-db", vertex_location="global", llm_router=None + ) + setup_data = {"model": setup_model} + + resolved = _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup_data, rewriter)) + + assert resolved == "gemini-live-2.5-flash", "an unresolved setup model logs the session as 'unknown'" + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_a_bare_alias_setup_model(): + """End to end through the relay: a bare alias must reach the logging object as a real model name. + + This is the call-site half of the fix. The helper tests above pass even if extraction moves back + before the rewrite, so this one drives the real websocket relay and asserts on what got logged, + which is the name the cost map is looked up by. An unbilled session logs ``unknown``. + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _build_vertex_live_setup_model_rewriter, + ) + + upstream_ws = RecordingUpstreamWebSocket() + setup_frame = json.dumps({"setup": {"model": "gemini-live-2.5-flash"}}) + websocket = _client_websocket( + AsyncMock( + side_effect=[ + {"type": "websocket.receive", "text": setup_frame}, + {"type": "websocket.disconnect"}, + ] + ) + ) + built = [] + real_logging = litellm.litellm_core_utils.litellm_logging.Logging + + def _capture(*args, **kwargs): + obj = real_logging(*args, **kwargs) + built.append(obj) + return obj + + with _patched_websocket_passthrough_environment(upstream_ws): + with patch("litellm.litellm_core_utils.litellm_logging.Logging", side_effect=_capture): + await websocket_passthrough_request( + websocket=websocket, + target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent", + custom_headers={"Authorization": "Bearer token"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/vertex_ai/live", + accept_websocket=False, + setup_model_rewriter=_build_vertex_live_setup_model_rewriter( + vertex_project="proj-db", vertex_location="global", llm_router=None + ), + ) + + assert built, "the relay should have built a logging object" + assert built[0].model == "gemini-live-2.5-flash", "a bare alias must not log as 'unknown'" + + +def test_vertex_live_setup_resolution_is_inert_without_a_rewriter(): + """Non-Live passthrough routes pass no rewriter, so the frame must be handed over untouched.""" + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _extract_model_from_vertex_ai_setup, + _resolved_vertex_live_setup, + ) + + setup_data = {"model": "projects/p/locations/global/publishers/google/models/gemini-live-2.5-flash"} + + assert _resolved_vertex_live_setup(setup_data, None) is setup_data + assert _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup_data, None)) == ( + "gemini-live-2.5-flash" + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("rcvd_close", [None, "abnormal", "no_status"]) async def test_websocket_passthrough_does_not_relay_unsendable_upstream_close(rcvd_close): From 6d71e3385b8c2ed0db61cfc1991de5614de15d83 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 22:12:46 -0400 Subject: [PATCH 016/107] fix(cost): carry grounding counters through the Responses usage bridge A realtime session's usage is rebuilt from its own response.done event, so a counter that does not survive the round trip is invisible to the cost path. Both directions copied a fixed allow-list, which meant a grounded Gemini Live session reported its query on the Usage object and then lost it before anything could bill it. Gemini reads the grounding counters off the input token details while Anthropic reads its own server_tool_use field, so carrying these two cannot move an Anthropic bill. Absent counters stay absent, so no provider starts paying a fee it did not incur. (cherry picked from commit ffd6c723e2a65dfff1860a913c34e41bc72ff11f) (cherry picked from commit 1590822f6893c57086b72965963d23f4f367b424) --- .../litellm_completion_transformation/transformation.py | 8 ++++++++ litellm/responses/utils.py | 6 ++++++ 2 files changed, 14 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b2d1a69e0d8..2b011abf0e2 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2667,6 +2667,14 @@ class LiteLLMCompletionResponsesConfig: if hasattr(prompt_details, "audio_tokens") and prompt_details.audio_tokens is not None: input_details_dict["audio_tokens"] = prompt_details.audio_tokens + # The cost path reads the grounding counters off the input details, and a realtime + # session's usage is rebuilt from its own response.done, so dropping them here bills + # no per-query grounding fee at all. + for counter in ("web_search_requests", "google_maps_grounding_requests"): + counter_value = getattr(prompt_details, counter, None) + if counter_value is not None: + input_details_dict[counter] = counter_value + cache_write_tokens = getattr(prompt_details, "cache_write_tokens", None) or getattr( prompt_details, "cache_creation_tokens", None ) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 540d492beec..7e0acd14b9c 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1137,6 +1137,12 @@ class ResponseAPILoggingUtils: text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), + web_search_requests=getattr( + response_api_usage.input_tokens_details, "web_search_requests", None + ), + google_maps_grounding_requests=getattr( + response_api_usage.input_tokens_details, "google_maps_grounding_requests", None + ), ) completion_tokens_details: CompletionTokensDetailsWrapper | None = None output_tokens_details: Final[OutputTokensDetails | None] = getattr( From 6c83484065ac1b6fbed79bb9d0acea97f50a0915 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Tue, 8 Sep 2026 00:59:28 -0400 Subject: [PATCH 017/107] chore(vertex-live): drop a comment that restated the helper's docstring The call site repeated _resolved_vertex_live_setup's own docstring almost verbatim, which is the duplication the repo's comment rule exists to prevent. --- litellm/proxy/pass_through_endpoints/pass_through_endpoints.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5a2f9c391a2..b66c295d1aa 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2330,9 +2330,6 @@ async def websocket_passthrough_request( setup_data, ) if isinstance(setup_data, dict) and "model" in setup_data: - # Resolve the alias first: a client may name a bare gateway model, - # which carries no "/models/" for the extractor to read, so reading - # the raw frame leaves the session logged as "unknown" and unbilled. extracted_model = _extract_model_from_vertex_ai_setup( _resolved_vertex_live_setup(setup_data, setup_model_rewriter) ) From 4e51ff8a6dbcc6a544891f8e64b887c8cd6a8ae2 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Wed, 9 Sep 2026 00:00:46 -0400 Subject: [PATCH 018/107] style(cost): collapse a getattr call that fits the line limit --- litellm/responses/utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 7e0acd14b9c..b2844834860 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1137,9 +1137,7 @@ class ResponseAPILoggingUtils: text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), - web_search_requests=getattr( - response_api_usage.input_tokens_details, "web_search_requests", None - ), + web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None), google_maps_grounding_requests=getattr( response_api_usage.input_tokens_details, "google_maps_grounding_requests", None ), From 228d87db6341a87bd9ae2bd1ebd76ae5e08e1a08 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Wed, 9 Sep 2026 00:19:12 -0400 Subject: [PATCH 019/107] test(gemini-live): assert grounding counters on the emitted response.done event Replaces a patch.object spy on a static method with a read of the public return value, which also covers the usage bridge the spy ran ahead of. --- .../test_gemini_realtime_transformation.py | 66 +++++++------------ 1 file changed, 22 insertions(+), 44 deletions(-) diff --git a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index ee984cc7e1f..3eb4a70ee15 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,7 +1,7 @@ import json from collections.abc import Mapping from typing import cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock import pytest @@ -9,8 +9,6 @@ import pytest import litellm from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig from litellm.types.llms.gemini import BidiGenerateContentServerMessage -from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents -from litellm.types.utils import Usage def test_gemini_realtime_transformation_session_created(): @@ -2205,40 +2203,22 @@ def _grounded_live_frame(grounding_metadata: Mapping[str, object] | None) -> Map } -def _usage_built_for_response_done(message: Mapping[str, object]) -> Usage: - """Capture the chat-completion Usage transform_response_done_event builds, before it is bridged. - - The Usage object is local to the method, so the bridge call is the only place it is observable. - """ +def _response_done_input_details(message: Mapping[str, object]) -> Mapping[str, object]: + """The ``input_tokens_details`` a ``response.done`` event carries, read off the emitted event.""" from typing import Final - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) - - captured: Final[list[Usage]] = [] # mutable-ok: a spy has to accumulate what it observes - original: Final = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage - - def _spy(usage: Usage) -> object: - captured.append(usage) - return original(usage) - config: Final = GeminiRealtimeConfig() - with patch.object( - LiteLLMCompletionResponsesConfig, - "_transform_chat_completion_usage_to_responses_usage", - staticmethod(_spy), - ): - config.transform_response_done_event( - message=cast( # cast-ok: a test fixture stands in for the server frame TypedDict - BidiGenerateContentServerMessage, message - ), - current_response_id="resp_grounding", - current_conversation_id="conv_grounding", - output_items=None, - ) - assert captured, "response.done must build a Usage object" - return captured[0] + event: Final = config.transform_response_done_event( + message=cast( # cast-ok: a test fixture stands in for the server frame TypedDict + BidiGenerateContentServerMessage, message + ), + current_response_id="resp_grounding", + current_conversation_id="conv_grounding", + output_items=None, + ) + usage: Final = event["response"]["usage"] + assert usage, "response.done must carry a usage object" + return usage.get("input_tokens_details") or {} def test_gemini_realtime_response_done_counts_web_grounding(): @@ -2247,12 +2227,10 @@ def test_gemini_realtime_response_done_counts_web_grounding(): Nothing read those frames on the realtime path, so web_search_requests stayed unset and the cost path's only trigger for Google's per-query grounding charge never fired. - Scope boundary, deliberate: this asserts the counter on the Usage object that response.done is - built from, not on the emitted event. The Responses usage bridge copies a fixed allow-list of - detail fields and drops the rest, so the counter does not reach response.done yet. Widening - that bridge is a separate change; do not read this test as proving end-to-end billing. + The counter is read off the emitted event, which is what the cost path is handed, so this covers + the grounding read and the usage bridge that carries it together """ - usage = _usage_built_for_response_done( + input_details = _response_done_input_details( _grounded_live_frame( { "webSearchQueries": ["who won the 2026 world cup final"], @@ -2261,13 +2239,13 @@ def test_gemini_realtime_response_done_counts_web_grounding(): ) ) - assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query" - assert usage.prompt_tokens_details.text_tokens == 19, "the modality breakdown must survive alongside it" + assert input_details.get("web_search_requests") == 1, "a grounded turn must report its query" + assert input_details.get("text_tokens") == 19, "the modality breakdown must survive alongside it" def test_gemini_realtime_response_done_reports_no_grounding_when_none_ran(): """The counter must stay unset on an ordinary turn, or every session pays a grounding fee.""" - usage = _usage_built_for_response_done(_grounded_live_frame(None)) + input_details = _response_done_input_details(_grounded_live_frame(None)) - assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None - assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None + assert input_details.get("web_search_requests") is None + assert input_details.get("google_maps_grounding_requests") is None From e205be80df1e23a22b4a5eee5f37bc9d52cb1ce1 Mon Sep 17 00:00:00 2001 From: oliver Date: Fri, 11 Sep 2026 07:41:13 +0000 Subject: [PATCH 020/107] fix(proxy): enforce custom_key_update policy on /key/regenerate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 47 ++++-- .../test_key_management_endpoints.py | 150 ++++++++++++++++++ 2 files changed, 184 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 749a940de0e..106e9b8027b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -294,6 +294,35 @@ def _custom_key_update_hook( return hooks.user_custom_key_update +async def _enforce_custom_key_update_policy( + hook: Callable[..., Awaitable[Mapping[str, object]]] | None, + data: UpdateKeyRequest, +) -> None: + if hook is None: + return + if not inspect.iscoroutinefunction(hook): + raise ValueError("user_custom_key_update must be a coroutine") + result: Final = await hook(data) + if not result.get("decision", True): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=result.get("message", "Authentication Failed - Custom Auth Rule"), + ) + + +def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None: + changed_fields: Final = MappingProxyType( + { + field: value + for field, value in data.model_dump(exclude_unset=True).items() + if field in UpdateKeyRequest.model_fields and field != "key" + } + ) + if not changed_fields: + return None + return UpdateKeyRequest(key=key, **changed_fields) + + class _LegacyDumpable(Protocol): def dict(self) -> Mapping[str, object]: ... @@ -3066,19 +3095,7 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) - # Custom key update hook - custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook( - proxy_server - ) - if custom_key_update_hook is not None: - if inspect.iscoroutinefunction(custom_key_update_hook): - result: Final = await custom_key_update_hook(data) - else: - raise ValueError("user_custom_key_update must be a coroutine") - decision: Final = result.get("decision", True) - message: Final = result.get("message", "Authentication Failed - Custom Auth Rule") - if not decision: - raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message) + await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=data) # Enforce upperbound key params on update (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) @@ -5069,6 +5086,7 @@ async def _execute_virtual_key_regeneration( proxy_logging_obj: ProxyLogging, ) -> GenerateKeyResponse: """Generate new token, update DB, invalidate cache, and return response.""" + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import hash_token # Mirror the /key/update ownership rebind guard. See helper docstring. @@ -5116,6 +5134,9 @@ async def _execute_virtual_key_regeneration( non_default_values = {} if data is not None: + update_request: Final = _regenerate_request_as_update_request(key=hashed_api_key, data=data) + if update_request is not None: + await _enforce_custom_key_update_policy(hook=_custom_key_update_hook(proxy_server), data=update_request) # Enforce upperbound key params on regenerate (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) non_default_values = await prepare_key_update_data(data=data, existing_key_row=key_in_db) 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 65cc23ea67f..fbe790dbed5 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 @@ -24,18 +24,21 @@ from litellm.proxy._types import ( LitellmUserRoles, Member, ProxyException, + RegenerateKeyRequest, ResetSpendRequest, UpdateKeyRequest, ) from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_org_key_limits, _check_project_key_limits, _check_team_key_limits, _common_key_generation_helper, _enforce_upperbound_key_params, + _execute_virtual_key_regeneration, _get_and_validate_existing_key, _list_key_helper, _persist_deleted_verification_tokens, @@ -12011,6 +12014,153 @@ async def test_execute_virtual_key_regeneration_allows_within_limit_duration(mon assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_rejects_when_custom_key_update_hook_denies(): + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="3000d") + mock_prisma_client = _make_regenerate_mock_prisma() + received_data: list[UpdateKeyRequest] = [] + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + received_data.append(data) + if data.duration and duration_in_seconds(data.duration) > duration_in_seconds("7d"): + return {"decision": False, "message": "duration must be <= 7d"} + return {"decision": True} + + with ( + patch( # test-quality-ok: deterministic token setup for policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ) as insert_deprecated_key, + patch( # test-quality-ok: cache eviction is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + with pytest.raises(HTTPException) as exc_info: + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == "duration must be <= 7d" + insert_deprecated_key.assert_not_awaited() + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + assert len(received_data) == 1 + assert received_data[0].key == "abc123" + assert received_data[0].duration == "3000d" + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_hook_approves(): + existing_key = _make_regenerate_existing_key() + data = RegenerateKeyRequest(duration="5d") + mock_prisma_client = _make_regenerate_mock_prisma() + received_data: list[UpdateKeyRequest] = [] + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + received_data.append(data) + if data.duration and duration_in_seconds(data.duration) > duration_in_seconds("7d"): + return {"decision": False, "message": "duration must be <= 7d"} + return {"decision": True} + + with ( + patch( # test-quality-ok: deterministic token setup for policy approval + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside policy approval + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: cache eviction is outside policy approval + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside policy approval + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + assert len(received_data) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()]) +async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_without_changes(data): + mock_prisma_client = _make_regenerate_mock_prisma() + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + raise AssertionError(f"custom key update hook called with {data}") + + with ( + patch( # test-quality-ok: deterministic token setup for unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: cache eviction is outside unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside unchanged request + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=_make_regenerate_existing_key(), + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + + @pytest.mark.asyncio async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token(): """ From 560e5da320bc00636057b3de3edd05d6a2da1745 Mon Sep 17 00:00:00 2001 From: oliver Date: Fri, 11 Sep 2026 08:02:18 +0000 Subject: [PATCH 021/107] fix(proxy): archive the old key only after regenerate validation passes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 18 +++++++---------- .../test_key_management_endpoints.py | 20 +++++++++++++++---- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 106e9b8027b..97e2a4f3be5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -5155,6 +5155,13 @@ async def _execute_virtual_key_regeneration( prisma_client=prisma_client, ) + await _persist_deleted_verification_tokens( + keys=[key_in_db], + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + # If grace period set, insert deprecated key so old key remains valid await _insert_deprecated_key( prisma_client=prisma_client, @@ -5453,17 +5460,6 @@ async def regenerate_key_fn( if litellm_changed_by is not None and not isinstance(litellm_changed_by, str): litellm_changed_by = None - # Save the old key record to deleted table before regeneration. - # This preserves key_alias and team_id metadata for historical spend records. - # If this fails, abort the regeneration to avoid permanently losing the - # old hash→metadata mapping. - await _persist_deleted_verification_tokens( - keys=[_key_in_db], - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=litellm_changed_by, - ) - return await _execute_virtual_key_regeneration( prisma_client=prisma_client, key_in_db=_key_in_db, 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 fbe790dbed5..6ec04815420 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 @@ -11935,6 +11935,10 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", new_callable=AsyncMock, ), + patch( # test-quality-ok: archival path is outside upperbound rejection + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + new_callable=AsyncMock, + ) as persist_deleted_verification_tokens, patch( "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", new_callable=AsyncMock, @@ -11955,6 +11959,7 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk assert exc_info.value.status_code == 400 assert "duration" in str(exc_info.value.detail) # Rejected regenerate must not reach the DB update. + persist_deleted_verification_tokens.assert_not_awaited() assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 @@ -12037,6 +12042,10 @@ async def test_execute_virtual_key_regeneration_rejects_when_custom_key_update_h "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", new_callable=AsyncMock, ) as insert_deprecated_key, + patch( # test-quality-ok: archival path is outside policy rejection + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + new_callable=AsyncMock, + ) as persist_deleted_verification_tokens, patch( # test-quality-ok: cache eviction is outside policy rejection "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", new_callable=AsyncMock, @@ -12063,6 +12072,7 @@ async def test_execute_virtual_key_regeneration_rejects_when_custom_key_update_h assert exc_info.value.status_code == 403 assert exc_info.value.detail == "duration must be <= 7d" insert_deprecated_key.assert_not_awaited() + persist_deleted_verification_tokens.assert_not_awaited() assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 assert len(received_data) == 1 assert received_data[0].key == "abc123" @@ -12092,6 +12102,10 @@ async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_ho "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", new_callable=AsyncMock, ), + patch( # test-quality-ok: verify archival follows policy approval + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + new_callable=AsyncMock, + ) as persist_deleted_verification_tokens, patch( # test-quality-ok: cache eviction is outside policy approval "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", new_callable=AsyncMock, @@ -12115,6 +12129,8 @@ async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_ho ) assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + persist_deleted_verification_tokens.assert_awaited_once() + assert persist_deleted_verification_tokens.call_args.kwargs["keys"] == [existing_key] assert len(received_data) == 1 @@ -13922,10 +13938,6 @@ async def test_regenerate_applies_normalized_mcp_object_permission(): "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_vector_stores_against_team", new_callable=AsyncMock, ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", - new_callable=AsyncMock, - ), patch( "litellm.proxy.management_endpoints.key_management_endpoints._execute_virtual_key_regeneration", execute_mock, From 577a4e94aa383bc48271f522a2f0e00804f749a1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:09:24 -0700 Subject: [PATCH 022/107] feat(proxy): unified custom_key_policy hook for key generate, update and regenerate Adds general_settings.custom_key_policy, one coroutine that receives the operation ("generate", "update", "regenerate"), the existing key row, the effective row as it will be written, and the raw request, and can deny with a 403. It runs after the request has been normalized and before the first DB write on /key/generate, /key/service-account/generate, /key/update, /key/bulk_update, /team/key/bulk_update and /key/{key}/regenerate. The two legacy hooks keep running unchanged on the raw request. --- .../key_management_endpoints.py | 160 +++++ litellm/proxy/proxy_server.py | 9 + .../key_management_endpoints.py | 17 + .../test_key_management_endpoints.py | 621 ++++++++++++++++++ .../test_db_overlay_remote_module_scrub.py | 1 + 5 files changed, 808 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ba43943d6de..a1ce1700f68 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -148,6 +148,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateKeyRequest, BulkUpdateKeyResponse, BulkUpdateTeamKeysRequest, + CustomKeyPolicyRequest, FailedKeyUpdate, KeySearchWhere, SuccessfulKeyUpdate, @@ -280,6 +281,7 @@ def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions: class _CustomKeyHooksModule(Protocol): user_custom_key_generate: Callable[..., Awaitable[Mapping[str, object]]] | None user_custom_key_update: Callable[..., Awaitable[Mapping[str, object]]] | None + user_custom_key_policy: Callable[..., Awaitable[Mapping[str, object]]] | None def _custom_key_generate_hook( @@ -294,6 +296,12 @@ def _custom_key_update_hook( return hooks.user_custom_key_update +def _custom_key_policy_hook( + hooks: _CustomKeyHooksModule, +) -> Callable[..., Awaitable[Mapping[str, object]]] | None: + return hooks.user_custom_key_policy + + async def _enforce_custom_key_update_policy( hook: Callable[..., Awaitable[Mapping[str, object]]] | None, data: UpdateKeyRequest, @@ -310,6 +318,113 @@ async def _enforce_custom_key_update_policy( ) +async def _enforce_custom_key_policy( + hook: Callable[..., Awaitable[Mapping[str, object]]] | None, + build_policy_request: Callable[[], CustomKeyPolicyRequest], +) -> None: + if hook is None: + return + if not inspect.iscoroutinefunction(hook): + raise ValueError("user_custom_key_policy must be a coroutine") + result: Final = await hook(build_policy_request()) + if not result.get("decision", True): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=result.get("message", "Authentication Failed - Custom Auth Rule"), + ) + + +_KEY_UPDATE_JSON_STRING_COLUMNS: Final = frozenset({"router_settings", "budget_limits"}) + +_KEY_METADATA_REQUEST_FIELDS: Final = frozenset( + (*LiteLLM_ManagementEndpoint_MetadataFields_Premium, *LiteLLM_ManagementEndpoint_MetadataFields) +) + + +def _decode_json_string_column(column: str, value: object) -> object: + if column in _KEY_UPDATE_JSON_STRING_COLUMNS and isinstance(value, str): + return json.loads(value) + return value + + +def _verification_token_from_row(row: Mapping[str, object]) -> LiteLLM_VerificationToken: + org_id: Final = row["organization_id"] if "organization_id" in row else row.get("org_id") + return LiteLLM_VerificationToken.model_validate(MappingProxyType({**row, "org_id": org_id})) + + +def _effective_key_after_update( + existing_key_row: LiteLLM_VerificationToken, + non_default_values: Mapping[str, object], +) -> LiteLLM_VerificationToken: + overlay: Final = MappingProxyType( + {column: _decode_json_string_column(column, value) for column, value in non_default_values.items()} + ) + return _verification_token_from_row(MappingProxyType({**existing_key_row.model_dump(), **overlay})) + + +def _update_policy_request( + operation: Literal["update", "regenerate"], + existing_key_row: LiteLLM_VerificationToken, + non_default_values: Mapping[str, object], + request: UpdateKeyRequest | RegenerateKeyRequest, +) -> CustomKeyPolicyRequest: + return CustomKeyPolicyRequest( + operation=operation, + existing_key=_verification_token_from_row(existing_key_row.model_dump()), + effective_key=_effective_key_after_update( + existing_key_row=existing_key_row, non_default_values=non_default_values + ), + request=request, + ) + + +def _generate_budget_windows( + budget_limits: Sequence[BudgetLimitEntry] | None, +) -> tuple[Mapping[str, object], ...] | None: + if budget_limits is None: + return None + return tuple( + MappingProxyType( + { + **window.model_dump(), + "reset_at": get_budget_reset_time(budget_duration=window.budget_duration).isoformat(), + } + ) + for window in budget_limits + ) + + +def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> LiteLLM_VerificationToken: + requested: Final = data.model_dump(exclude_unset=True, exclude_none=True) + metadata_fields: Final = MappingProxyType( + {field: value for field, value in requested.items() if field in _KEY_METADATA_REQUEST_FIELDS} + ) + column_fields: Final = MappingProxyType( + {field: value for field, value in requested.items() if field not in _KEY_METADATA_REQUEST_FIELDS} + ) + request_metadata: Final = data.metadata or MappingProxyType({}) + folded_metadata: Final = {**request_metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict + columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates data_json in place + expires: Final = ( + now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None + ) + budget_reset_at: Final = ( + get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None + ) + return _verification_token_from_row( + MappingProxyType( + { + **columns, + "metadata": encrypt_callback_vars(folded_metadata), + "expires": expires, + "budget_reset_at": budget_reset_at, + "budget_limits": _generate_budget_windows(data.budget_limits), + "object_permission": None, + } + ) + ) + + def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None: changed_fields: Final = MappingProxyType( { @@ -1016,6 +1131,7 @@ async def _common_key_generation_helper( litellm_changed_by: str | None, team_table: LiteLLM_TeamTableCachedObj | None, ) -> GenerateKeyResponse: + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, llm_router, @@ -1164,6 +1280,16 @@ async def _common_key_generation_helper( "litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - %s", e ) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: CustomKeyPolicyRequest( + operation="generate", + existing_key=None, + effective_key=_effective_key_for_generate(data=data, now=datetime.now(timezone.utc)), + request=data, + ), + ) + # TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable _budget_id = data.budget_id if prisma_client is not None and data.soft_budget is not None: @@ -2496,6 +2622,7 @@ async def _process_single_key_update( llm_router: Router | None, user_custom_key_update: Callable | None = None, existing_key_row: LiteLLM_VerificationToken | None = None, + user_custom_key_policy: Callable[..., Awaitable[Mapping[str, object]]] | None = None, ) -> dict[str, object]: """ Process a single key update with all validations and checks. @@ -2606,6 +2733,16 @@ async def _process_single_key_update( # Prepare update data non_default_values = await prepare_key_update_data(data=update_key_request, existing_key_row=existing_key_row) + await _enforce_custom_key_policy( + hook=user_custom_key_policy, + build_policy_request=lambda: _update_policy_request( + operation="update", + existing_key_row=existing_key_row, + non_default_values=non_default_values, + request=update_key_request, + ), + ) + # Update key in database if prisma_client is None: raise HTTPException( @@ -3112,6 +3249,16 @@ async def update_key_fn( _enforce_upperbound_key_params(data, fill_defaults=False) non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: _update_policy_request( + operation="update", + existing_key_row=existing_key_row, + non_default_values=non_default_values, + request=data, + ), + ) + # Only validate key_alias format if it's actually being changed new_key_alias: Final = non_default_values.get("key_alias", None) if new_key_alias != existing_key_row.key_alias: @@ -3280,6 +3427,7 @@ async def bulk_update_keys( ) custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + custom_key_policy_hook: Final = _custom_key_policy_hook(proxy_server) if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( @@ -3327,6 +3475,7 @@ async def bulk_update_keys( proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, user_custom_key_update=custom_key_update_hook, + user_custom_key_policy=custom_key_policy_hook, ) successful_updates.append( @@ -3444,6 +3593,7 @@ async def bulk_update_team_keys( ) custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + custom_key_policy_hook: Final = _custom_key_policy_hook(proxy_server) if prisma_client is None: raise HTTPException( @@ -3574,6 +3724,7 @@ async def bulk_update_team_keys( proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, user_custom_key_update=custom_key_update_hook, + user_custom_key_policy=custom_key_policy_hook, existing_key_row=existing_by_token[db_token], ) @@ -5156,6 +5307,15 @@ async def _execute_virtual_key_regeneration( if new_key_alias != key_in_db.key_alias: _validate_key_alias_format(key_alias=new_key_alias) verbose_proxy_logger.debug("non_default_values: %s", non_default_values) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: _update_policy_request( + operation="regenerate", + existing_key_row=key_in_db, + non_default_values=non_default_values, + request=data if data is not None else RegenerateKeyRequest(), + ), + ) update_data.update(non_default_values) jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 606a590c24b..9d258f6fd92 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -924,6 +924,7 @@ def cleanup_router_config_variables(): user_custom_auth_path, \ user_custom_key_generate, \ user_custom_key_update, \ + user_custom_key_policy, \ user_custom_sso, \ user_custom_ui_sso_sign_in_handler, \ use_background_health_checks, \ @@ -941,6 +942,7 @@ def cleanup_router_config_variables(): user_custom_auth_path = None user_custom_key_generate = None user_custom_key_update = None + user_custom_key_policy = None TEAM_METADATA_VALIDATOR_REGISTRY.set(None) TEAM_METADATA_SCHEMA_REGISTRY.set(()) user_custom_sso = None @@ -2365,6 +2367,7 @@ user_custom_key_generate = None _pkce_no_redis_warning_emitted: bool = False _cp_no_redis_warning_emitted: bool = False user_custom_key_update = None +user_custom_key_policy = None user_custom_sso = None user_custom_ui_sso_sign_in_handler = None use_background_health_checks = None @@ -4250,6 +4253,7 @@ _DB_OVERLAY_REMOTE_MODULE_STR_FIELDS: Final[dict[str, tuple[str, ...]]] = { "custom_auth", "custom_key_generate", "custom_key_update", + "custom_key_policy", "custom_team_metadata_validate", "custom_sso", "custom_ui_sso_sign_in_handler", @@ -5403,6 +5407,7 @@ class ProxyConfig: user_custom_auth_path, \ user_custom_key_generate, \ user_custom_key_update, \ + user_custom_key_policy, \ user_custom_sso, \ user_custom_ui_sso_sign_in_handler, \ use_background_health_checks, \ @@ -5940,6 +5945,10 @@ class ProxyConfig: if custom_key_update is not None: user_custom_key_update = get_instance_fn(value=custom_key_update, config_file_path=config_file_path) + custom_key_policy: Final = general_settings.get("custom_key_policy", None) + if custom_key_policy is not None: + user_custom_key_policy = get_instance_fn(value=custom_key_policy, config_file_path=config_file_path) + custom_team_metadata_validate: Final = general_settings.get("custom_team_metadata_validate", None) TEAM_METADATA_VALIDATOR_REGISTRY.set( get_instance_fn(value=custom_team_metadata_validate, config_file_path=config_file_path) diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 9fb5bea81e3..35abb8fb718 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -4,6 +4,9 @@ from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, model_validator from typing_extensions import ReadOnly, TypedDict +from litellm.models.verification_token import LiteLLM_VerificationToken +from litellm.proxy._types import GenerateKeyRequest, RegenerateKeyRequest, UpdateKeyRequest +from litellm.types.llms.base import LiteLLMPydanticObjectBase from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains @@ -123,3 +126,17 @@ class BulkUpdateTeamKeysRequest(BaseModel): if not has_key_ids and not self.all_keys_in_team: raise ValueError("Must provide either `key_ids` (non-empty) or `all_keys_in_team=True`.") return self + + +CustomKeyPolicyOperation = Literal["generate", "update", "regenerate"] + + +class CustomKeyPolicyRequest(LiteLLMPydanticObjectBase): + """What `general_settings.custom_key_policy` receives: the operation, the key row as it will be written, and the raw request.""" + + model_config = ConfigDict(protected_namespaces=(), frozen=True) + + operation: CustomKeyPolicyOperation + existing_key: LiteLLM_VerificationToken | None + effective_key: LiteLLM_VerificationToken + request: GenerateKeyRequest | UpdateKeyRequest | RegenerateKeyRequest 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 000a9dc2f1a..2b8c62523b9 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 @@ -1,3 +1,4 @@ +from contextlib import ExitStack from typing import Final import json from datetime import datetime, timedelta, timezone @@ -22,6 +23,7 @@ from litellm.proxy._types import ( LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LiteLLM_VerificationToken, + LiteLLMKeyType, LitellmUserRoles, Member, ProxyException, @@ -38,6 +40,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_project_key_limits, _check_team_key_limits, _common_key_generation_helper, + _effective_key_after_update, + _effective_key_for_generate, + _enforce_custom_key_policy, _enforce_upperbound_key_params, _execute_virtual_key_regeneration, _get_and_validate_existing_key, @@ -64,6 +69,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( validate_key_team_change, ) from litellm.proxy.proxy_server import app +from litellm.types.proxy.management_endpoints.key_management_endpoints import CustomKeyPolicyRequest client = TestClient(app) @@ -12178,6 +12184,586 @@ async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_wit assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 +_POLICY_DENIAL_MESSAGE = "key duration must be 7d or less" +_POLICY_HASHED_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" +_POLICY_GENERATED_KEY = {"key": "sk-test-key", "expires": None, "user_id": "test-user", "team_id": None} + + +def _seven_day_policy(received: list[CustomKeyPolicyRequest]): + async def policy(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + expires = policy_request.effective_key.expires + if isinstance(expires, datetime) and expires > datetime.now(timezone.utc) + timedelta(days=7): + return {"decision": False, "message": _POLICY_DENIAL_MESSAGE} + return {"decision": True} + + return policy + + +def _assert_expires_in(effective_key: LiteLLM_VerificationToken, duration: str) -> None: + expires = effective_key.expires + assert isinstance(expires, datetime) + assert expires.tzinfo is not None + expected = datetime.now(timezone.utc) + timedelta(seconds=duration_in_seconds(duration=duration)) + assert abs((expires - expected).total_seconds()) < 60 + + +def _regenerate_policy_mocks(policy, insert_deprecated_key: AsyncMock, persist: AsyncMock) -> ExitStack: + stack = ExitStack() + stack.enter_context( + patch( # test-quality-ok: deterministic token setup for the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ) + ) + stack.enter_context( + patch( # test-quality-ok: grace-period write must not run on a denied regenerate + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + insert_deprecated_key, + ) + ) + stack.enter_context( + patch( # test-quality-ok: archival write must not run on a denied regenerate + "litellm.proxy.management_endpoints.key_management_endpoints._persist_deleted_verification_tokens", + persist, + ) + ) + stack.enter_context( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ) + ) + stack.enter_context( + patch( # test-quality-ok: rotation callback is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ) + ) + stack.enter_context( + patch("litellm.proxy.proxy_server.user_custom_key_policy", policy) # test-quality-ok: inject policy hook + ) + return stack + + +async def _regenerate_under_policy(mock_prisma_client, existing_key, data): + return await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_regenerate_rejects_when_custom_key_policy_denies_the_effective_expiry(): + existing_key = _make_regenerate_existing_key() + mock_prisma_client = _make_regenerate_mock_prisma() + received: list[CustomKeyPolicyRequest] = [] + insert_deprecated_key = AsyncMock() + persist = AsyncMock() + + with _regenerate_policy_mocks(_seven_day_policy(received), insert_deprecated_key, persist): + with pytest.raises(HTTPException) as exc_info: + await _regenerate_under_policy(mock_prisma_client, existing_key, RegenerateKeyRequest(duration="3000d")) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE + insert_deprecated_key.assert_not_awaited() + persist.assert_not_awaited() + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + assert [policy_request.operation for policy_request in received] == ["regenerate"] + assert received[0].existing_key is not None + assert received[0].existing_key.token == "abc123" + assert isinstance(received[0].request, RegenerateKeyRequest) + assert received[0].request.duration == "3000d" + _assert_expires_in(received[0].effective_key, "3000d") + + +@pytest.mark.asyncio +async def test_regenerate_within_custom_key_policy_rotates_the_key(): + existing_key = _make_regenerate_existing_key() + mock_prisma_client = _make_regenerate_mock_prisma() + received: list[CustomKeyPolicyRequest] = [] + persist = AsyncMock() + + with _regenerate_policy_mocks(_seven_day_policy(received), AsyncMock(), persist): + await _regenerate_under_policy(mock_prisma_client, existing_key, RegenerateKeyRequest(duration="5d")) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + persist.assert_awaited_once() + assert persist.call_args.kwargs["keys"] == [existing_key] + assert [policy_request.operation for policy_request in received] == ["regenerate"] + _assert_expires_in(received[0].effective_key, "5d") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()]) +async def test_regenerate_without_changes_still_runs_custom_key_policy(data): + existing_key = _make_regenerate_existing_key() + mock_prisma_client = _make_regenerate_mock_prisma() + received: list[CustomKeyPolicyRequest] = [] + + async def freeze_rotation(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + return {"decision": False, "message": "key rotation is frozen"} + + with _regenerate_policy_mocks(freeze_rotation, AsyncMock(), AsyncMock()): + with pytest.raises(HTTPException) as exc_info: + await _regenerate_under_policy(mock_prisma_client, existing_key, data) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == "key rotation is frozen" + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + assert [policy_request.operation for policy_request in received] == ["regenerate"] + assert received[0].existing_key == existing_key + assert received[0].effective_key == existing_key + + +def _policy_existing_team_key() -> LiteLLM_VerificationToken: + return LiteLLM_VerificationToken( + token=_POLICY_HASHED_TOKEN, user_id="test-user", team_id="team-a", max_budget=200.0 + ) + + +def _setup_update_key_fn_policy_mocks(monkeypatch, existing_key: LiteLLM_VerificationToken) -> AsyncMock: + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key) + mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(return_value=None) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {"max_budget": 50.0, "team_id": "team-a"}}) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", AsyncMock(return_value=None) + ) + return mock_prisma_client + + +def _assert_update_policy_request(policy_request: CustomKeyPolicyRequest, request: UpdateKeyRequest) -> None: + assert policy_request.operation == "update" + assert policy_request.request is request + assert policy_request.existing_key is not None + assert policy_request.existing_key.max_budget == 200.0 + assert policy_request.effective_key.team_id == "team-a" + assert policy_request.effective_key.user_id == "test-user" + assert policy_request.effective_key.max_budget == 50.0 + _assert_expires_in(policy_request.effective_key, request.duration or "") + + +@pytest.mark.asyncio +async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) + received: list[CustomKeyPolicyRequest] = [] + policy = _seven_day_policy(received) + data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0) + + with ( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_policy", policy), # test-quality-ok: inject policy hook + ): + await update_key_fn( + request=MagicMock(), + data=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + mock_prisma_client.update_data.assert_awaited_once() + assert len(received) == 1 + _assert_update_policy_request(received[0], data) + + +@pytest.mark.asyncio +async def test_update_key_fn_rejects_when_custom_key_policy_denies(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) + received: list[CustomKeyPolicyRequest] = [] + policy = _seven_day_policy(received) + + with patch("litellm.proxy.proxy_server.user_custom_key_policy", policy): # test-quality-ok: inject policy hook + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="3000d", max_budget=50.0), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + assert str(exc_info.value.code) == "403" + assert exc_info.value.message == _POLICY_DENIAL_MESSAGE + mock_prisma_client.update_data.assert_not_awaited() + assert [policy_request.operation for policy_request in received] == ["update"] + _assert_expires_in(received[0].effective_key, "3000d") + + +async def _process_single_key_update_under_policy(prisma_client: AsyncMock, data: UpdateKeyRequest, policy): + with ( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: update callback is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", + new_callable=AsyncMock, + ), + ): + return await _process_single_key_update( + update_key_request=data, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + prisma_client=prisma_client, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + llm_router=None, + existing_key_row=_policy_existing_team_key(), + user_custom_key_policy=policy, + ) + + +@pytest.mark.asyncio +async def test_process_single_key_update_runs_custom_key_policy_on_the_effective_row(): + mock_prisma_client = AsyncMock() + updated_row = MagicMock() + updated_row.model_dump.return_value = {"max_budget": 50.0, "team_id": "team-a"} + mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row}) + received: list[CustomKeyPolicyRequest] = [] + data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0) + + result = await _process_single_key_update_under_policy(mock_prisma_client, data, _seven_day_policy(received)) + + assert result["max_budget"] == 50.0 + mock_prisma_client.update_data.assert_awaited_once() + assert len(received) == 1 + _assert_update_policy_request(received[0], data) + + +@pytest.mark.asyncio +async def test_process_single_key_update_rejects_when_custom_key_policy_denies(): + mock_prisma_client = AsyncMock() + mock_prisma_client.update_data = AsyncMock() + received: list[CustomKeyPolicyRequest] = [] + data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="3000d", max_budget=50.0) + + with pytest.raises(HTTPException) as exc_info: + await _process_single_key_update_under_policy(mock_prisma_client, data, _seven_day_policy(received)) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE + mock_prisma_client.update_data.assert_not_awaited() + assert [policy_request.operation for policy_request in received] == ["update"] + + +@pytest.mark.asyncio +async def test_bulk_update_keys_runs_custom_key_policy_per_key(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequest, + BulkUpdateKeyRequestItem, + ) + + existing_keys = [ + LiteLLM_VerificationToken(token="test-key-1", user_id="user-123", max_budget=None), + LiteLLM_VerificationToken(token="test-key-2", user_id="user-123", max_budget=50.0), + ] + updated_row = MagicMock() + updated_row.model_dump.return_value = {"user_id": "user-123", "max_budget": 100.0} + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=existing_keys) + mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row}) + mock_prisma_client.get_data = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + received: list[CustomKeyPolicyRequest] = [] + + async def cap_max_budget(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + max_budget = policy_request.effective_key.max_budget + if max_budget is not None and max_budget > 100: + return {"decision": False, "message": "max_budget must be 100 or less"} + return {"decision": True} + + monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", cap_max_budget) + + with ( + patch( # test-quality-ok: cache eviction is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: update callback is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", + new_callable=AsyncMock, + ), + ): + response = await bulk_update_keys( + data=BulkUpdateKeyRequest( + keys=[ + BulkUpdateKeyRequestItem(key="test-key-1", max_budget=100.0), + BulkUpdateKeyRequestItem(key="test-key-2", max_budget=500.0), + ] + ), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + assert [update.key for update in response.successful_updates] == ["test-key-1"] + assert [(failed.key, failed.failed_reason) for failed in response.failed_updates] == [ + ("test-key-2", "max_budget must be 100 or less") + ] + assert mock_prisma_client.update_data.await_count == 1 + assert [policy_request.operation for policy_request in received] == ["update", "update"] + assert [policy_request.effective_key.max_budget for policy_request in received] == [100.0, 500.0] + assert [ + policy_request.existing_key.max_budget if policy_request.existing_key is not None else "missing" + for policy_request in received + ] == [None, 50.0] + + +def _policy_generate_prisma() -> MagicMock: + mock_prisma = MagicMock() + mock_prisma.db.litellm_budgettable.create = AsyncMock(return_value=MagicMock(budget_id="budget-1")) + mock_prisma.jsonify_object = MagicMock(side_effect=lambda data: json.loads(data) if isinstance(data, str) else data) + return mock_prisma + + +def _generate_policy_mocks(mock_prisma: MagicMock, generate_key_helper: AsyncMock, policy) -> ExitStack: + stack = ExitStack() + stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)) # test-quality-ok: fake DB + stack.enter_context(patch("litellm.proxy.proxy_server.llm_router", None)) # test-quality-ok: no router in test + stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True)) # test-quality-ok: premium fields + stack.enter_context(patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")) # test-quality-ok: admin + stack.enter_context(patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())) # test-quality-ok: cache + stack.enter_context( + patch( # test-quality-ok: the key write must not run on a denied generate + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + generate_key_helper, + ) + ) + stack.enter_context( + patch("litellm.proxy.proxy_server.user_custom_key_policy", policy) # test-quality-ok: inject policy hook + ) + return stack + + +def _generate_request(duration: str, organization_id: str | None) -> GenerateKeyRequest: + return GenerateKeyRequest( + duration=duration, + organization_id=organization_id, + guardrails=["g1"], + tags=["t1"], + soft_budget=10.0, + max_budget=20.0, + ) + + +def _assert_generate_policy_request( + policy_request: CustomKeyPolicyRequest, duration: str, organization_id: str | None +) -> None: + assert policy_request.operation == "generate" + assert policy_request.existing_key is None + assert policy_request.effective_key.org_id == organization_id + assert policy_request.effective_key.max_budget == 20.0 + assert policy_request.effective_key.metadata["guardrails"] == ["g1"] + assert policy_request.effective_key.metadata["tags"] == ["t1"] + _assert_expires_in(policy_request.effective_key, duration) + + +@pytest.mark.asyncio +async def test_generate_key_rejects_when_custom_key_policy_denies_before_any_write(): + mock_prisma = _policy_generate_prisma() + generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY) + received: list[CustomKeyPolicyRequest] = [] + data = _generate_request("3000d", organization_id="org-1") + + with _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)): + with pytest.raises(ProxyException) as exc_info: + await generate_key_fn( + data=data, user_api_key_dict=_make_regenerate_user_api_key_dict(), litellm_changed_by=None + ) + + assert str(exc_info.value.code) == "403" + assert exc_info.value.message == _POLICY_DENIAL_MESSAGE + mock_prisma.db.litellm_budgettable.create.assert_not_awaited() + generate_key_helper.assert_not_awaited() + assert len(received) == 1 + _assert_generate_policy_request(received[0], "3000d", organization_id="org-1") + assert received[0].request is data + assert data.duration == "3000d" + assert data.guardrails == ["g1"] + assert data.tags == ["t1"] + assert data.organization_id == "org-1" + + +@pytest.mark.asyncio +async def test_generate_key_within_custom_key_policy_creates_the_key(): + mock_prisma = _policy_generate_prisma() + generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY) + received: list[CustomKeyPolicyRequest] = [] + data = _generate_request("5d", organization_id=None) + + with _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)): + await generate_key_fn( + data=data, user_api_key_dict=_make_regenerate_user_api_key_dict(), litellm_changed_by=None + ) + + mock_prisma.db.litellm_budgettable.create.assert_awaited_once() + generate_key_helper.assert_awaited_once() + assert len(received) == 1 + _assert_generate_policy_request(received[0], "5d", organization_id=None) + assert received[0].request is data + + +@pytest.mark.asyncio +async def test_service_account_generate_rejects_when_custom_key_policy_denies(): + from litellm.proxy.management_endpoints.key_management_endpoints import generate_service_account_key_fn + + mock_prisma = _policy_generate_prisma() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=MagicMock()) + generate_key_helper = AsyncMock(return_value=_POLICY_GENERATED_KEY) + received: list[CustomKeyPolicyRequest] = [] + + with ( + _generate_policy_mocks(mock_prisma, generate_key_helper, _seven_day_policy(received)), + patch( # test-quality-ok: team lookup is outside the policy path + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=None, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await generate_service_account_key_fn( + data=GenerateKeyRequest(team_id="team-1", duration="3000d"), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == _POLICY_DENIAL_MESSAGE + generate_key_helper.assert_not_awaited() + mock_prisma.db.litellm_budgettable.create.assert_not_awaited() + assert [policy_request.operation for policy_request in received] == ["generate"] + assert received[0].existing_key is None + assert received[0].effective_key.team_id == "team-1" + assert received[0].effective_key.user_id is None + _assert_expires_in(received[0].effective_key, "3000d") + + +@pytest.mark.asyncio +async def test_effective_key_after_update_decodes_json_string_columns_and_keeps_omitted_fields(): + existing_key = LiteLLM_VerificationToken(token="tok", user_id="u1", team_id="team-a") + non_default_values = await prepare_key_update_data( + data=UpdateKeyRequest( + key="tok", router_settings={"num_retries": 3}, budget_limits=[{"budget_duration": "1d", "max_budget": 2.0}] + ), + existing_key_row=existing_key, + ) + assert isinstance(non_default_values["router_settings"], str) + assert isinstance(non_default_values["budget_limits"], str) + + effective_key = _effective_key_after_update(existing_key_row=existing_key, non_default_values=non_default_values) + + assert effective_key.router_settings == {"num_retries": 3} + assert effective_key.budget_limits is not None + assert effective_key.budget_limits[0]["max_budget"] == 2.0 + assert effective_key.budget_limits[0]["budget_duration"] == "1d" + assert effective_key.budget_limits[0]["reset_at"] is not None + assert effective_key.team_id == "team-a" + assert effective_key.user_id == "u1" + + +@pytest.mark.asyncio +async def test_effective_key_after_update_clears_expiry_for_a_minus_one_duration(): + existing_key = LiteLLM_VerificationToken(token="tok", expires=datetime(2027, 1, 1, tzinfo=timezone.utc)) + non_default_values = await prepare_key_update_data( + data=UpdateKeyRequest(key="tok", duration="-1"), existing_key_row=existing_key + ) + + effective_key = _effective_key_after_update(existing_key_row=existing_key, non_default_values=non_default_values) + + assert effective_key.expires is None + + +def test_effective_key_for_generate_reflects_the_processed_request_without_mutating_it(): + now = datetime(2026, 1, 1, tzinfo=timezone.utc) + data = GenerateKeyRequest( + duration="5d", + organization_id="org-1", + metadata={"a": 1}, + guardrails=["g1"], + tags=["t1"], + budget_duration="1d", + max_budget=3.0, + key_type=LiteLLMKeyType.LLM_API, + ) + + effective_key = _effective_key_for_generate(data=data, now=now) + + assert effective_key.expires == now + timedelta(days=5) + assert effective_key.org_id == "org-1" + assert effective_key.metadata == {"a": 1, "guardrails": ["g1"], "tags": ["t1"]} + assert effective_key.max_budget == 3.0 + assert effective_key.budget_duration == "1d" + assert effective_key.budget_reset_at is not None + assert effective_key.key_type == "llm_api" + assert effective_key.allowed_routes == ["llm_api_routes"] + assert data.metadata == {"a": 1} + assert data.guardrails == ["g1"] + assert data.tags == ["t1"] + assert data.duration == "5d" + + +def test_effective_key_for_generate_without_duration_never_expires(): + effective_key = _effective_key_for_generate( + data=GenerateKeyRequest(), now=datetime(2026, 1, 1, tzinfo=timezone.utc) + ) + + assert effective_key.expires is None + assert effective_key.budget_reset_at is None + assert effective_key.key_type == "default" + + +def _policy_request_for_generate() -> CustomKeyPolicyRequest: + return CustomKeyPolicyRequest( + operation="generate", + existing_key=None, + effective_key=LiteLLM_VerificationToken(token="tok"), + request=GenerateKeyRequest(), + ) + + +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_rejects_a_sync_hook(): + def sync_hook(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + return {"decision": True} + + with pytest.raises(ValueError, match="user_custom_key_policy must be a coroutine"): + await _enforce_custom_key_policy(hook=sync_hook, build_policy_request=_policy_request_for_generate) + + +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_uses_the_default_denial_message(): + async def deny(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + return {"decision": False} + + with pytest.raises(HTTPException) as exc_info: + await _enforce_custom_key_policy(hook=deny, build_policy_request=_policy_request_for_generate) + + assert exc_info.value.status_code == 403 + assert exc_info.value.detail == "Authentication Failed - Custom Auth Rule" + + @pytest.mark.asyncio async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token(): """ @@ -18348,3 +18934,38 @@ async def test_key_creator_cannot_detach_project_without_admin_access(): ) assert exc.value.status_code == 403 assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_bulk_update_team_keys_runs_custom_key_policy_per_key(monkeypatch): + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateTeamKeysRequest, + KeyUpdateFields, + ) + + keys = [_make_team_key("tok-a"), _make_team_key("tok-b")] + mock = _setup_team_keys_mocks( + monkeypatch, find_many=keys, update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}) + ) + received: list[CustomKeyPolicyRequest] = [] + + async def freeze_tok_b(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + if policy_request.existing_key is not None and policy_request.existing_key.token == "tok-b": + return {"decision": False, "message": "tok-b is frozen"} + return {"decision": True} + + monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", freeze_tok_b) + + response = await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", key_ids=["tok-a", "tok-b"], update_fields=KeyUpdateFields(max_budget=50.0) + ) + ) + + assert [update.key for update in response.successful_updates] == ["tok-a"] + assert [(failed.key, failed.failed_reason) for failed in response.failed_updates] == [("tok-b", "tok-b is frozen")] + mock.update_data.assert_awaited_once() + assert [policy_request.operation for policy_request in received] == ["update", "update"] + assert [policy_request.effective_key.max_budget for policy_request in received] == [50.0, 50.0] + assert [policy_request.effective_key.team_id for policy_request in received] == ["team-abc", "team-abc"] diff --git a/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py b/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py index 100ba653f3a..0072997a0d9 100644 --- a/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py +++ b/tests/test_litellm/proxy/types_utils/test_db_overlay_remote_module_scrub.py @@ -43,6 +43,7 @@ def test_litellm_settings_callback_list_strips_remote_urls(field): "custom_auth", "custom_key_generate", "custom_key_update", + "custom_key_policy", "custom_sso", "custom_ui_sso_sign_in_handler", ], From 83594427fca6ca448426a7fd219ce68a832ba26c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:10:10 -0700 Subject: [PATCH 023/107] fix(vertex-live): price each grounded turn's query fee on the /vertex_ai/live passthrough --- ...tex_ai_live_passthrough_logging_handler.py | 130 +++++++++++++++--- .../test_vertex_ai_live_passthrough.py | 63 ++++++++- 2 files changed, 174 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 0e2eb60704d..0ac04654e30 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -7,11 +7,12 @@ Supports different modalities: text, audio, video, and web search. from collections.abc import Mapping, Sequence from datetime import datetime -from itertools import chain +from itertools import chain, pairwise from types import MappingProxyType -from typing import Final +from typing import Final, Literal, TypeAlias from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( BasePassthroughLoggingHandler, ) @@ -20,6 +21,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrou ) from litellm.types.utils import ( CompletionTokensDetailsWrapper, + CostBreakdown, LlmProviders, ModelResponse, PromptTokensDetailsWrapper, @@ -60,6 +62,35 @@ def _grounding_metadata(websocket_messages: Sequence[object]) -> tuple[Mapping[s ) +def _turns(websocket_messages: Sequence[object]) -> tuple[tuple[object, ...], ...]: + """Split a session at every ``usageMetadata`` frame; frames after the last one never got their usage.""" + closes: Final = tuple( + index + 1 + for index, message in enumerate(websocket_messages) + if isinstance(message, Mapping) and isinstance(message.get("usageMetadata"), dict) + ) + return tuple(tuple(websocket_messages[start:end]) for start, end in pairwise((0, *closes))) + + +_SummedField: TypeAlias = Literal[ + "input_cost", + "output_cost", + "tool_usage_cost", + "cache_read_cost", + "cache_creation_cost", + "reasoning_cost", + "original_cost", + "discount_amount", + "margin_fixed_amount", + "margin_total_amount", +] + + +def _summed(breakdowns: Sequence[CostBreakdown], field: _SummedField) -> float | None: + values: Final = tuple(value for breakdown in breakdowns if (value := breakdown.get(field)) is not None) + return sum(values) if values else None + + class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough. @@ -141,7 +172,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): @staticmethod def _extract_usage_metadata_from_websocket_messages( - websocket_messages: list[dict], + websocket_messages: Sequence[object], ) -> dict | None: """ Extract and aggregate usage metadata from a list of WebSocket messages. @@ -158,9 +189,11 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): Dictionary containing aggregated usage metadata, or None if not found """ snapshots: Final = tuple( - message["usageMetadata"] + metadata for message in websocket_messages - if isinstance(message, dict) and isinstance(message.get("usageMetadata"), dict) + if isinstance(message, Mapping) + for metadata in (message.get("usageMetadata"),) + if isinstance(metadata, dict) ) if not snapshots: @@ -247,10 +280,72 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): ) return usage + def _session_usage(self, websocket_messages: Sequence[object], model: str) -> Usage | None: + usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages) + if usage_metadata is None: + return None + return self._create_usage_object_from_metadata( + usage_metadata=usage_metadata, + grounding_metadata=_grounding_metadata(websocket_messages), + model=model, + ) + + def _turn_cost( + self, + turn: Sequence[object], + model: str, + logging_obj: LiteLLMLoggingObj, + ) -> tuple[float, CostBreakdown] | None: + usage: Final = self._session_usage(turn, model) + if usage is None: + return None + cost: Final = logging_obj._response_cost_calculator( # pyright: ignore[reportPrivateUsage] # the call's own calculator keeps custom pricing and the deployment's region in step with the spend row + result=ModelResponse(model=model, usage=usage), + litellm_model_name=model, + ) + if cost is None: + return None + breakdown: Final = logging_obj.cost_breakdown + return None if breakdown is None else (cost, breakdown) + + def _session_cost( + self, + websocket_messages: Sequence[object], + model: str, + logging_obj: LiteLLMLoggingObj, + ) -> float | None: + """Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice.""" + turn_costs: Final = tuple(self._turn_cost(turn, model, logging_obj) for turn in _turns(websocket_messages)) + priced: Final = tuple(turn_cost for turn_cost in turn_costs if turn_cost is not None) + if not priced or len(priced) != len(turn_costs): + return None + breakdowns: Final = tuple(breakdown for _, breakdown in priced) + first: Final = breakdowns[0] + total_cost: Final = sum(cost for cost, _ in priced) + logging_obj.set_cost_breakdown( + input_cost=_summed(breakdowns, "input_cost") or 0.0, + output_cost=_summed(breakdowns, "output_cost") or 0.0, + total_cost=total_cost, + cost_for_built_in_tools_cost_usd_dollar=_summed(breakdowns, "tool_usage_cost") or 0.0, + original_cost=_summed(breakdowns, "original_cost"), + discount_percent=first.get("discount_percent"), + discount_amount=_summed(breakdowns, "discount_amount"), + margin_percent=first.get("margin_percent"), + margin_fixed_amount=_summed(breakdowns, "margin_fixed_amount"), + margin_total_amount=_summed(breakdowns, "margin_total_amount"), + cache_read_cost=_summed(breakdowns, "cache_read_cost"), + cache_creation_cost=_summed(breakdowns, "cache_creation_cost"), + reasoning_cost=_summed(breakdowns, "reasoning_cost"), + service_tier=first.get("service_tier"), + data_residency=first.get("data_residency"), + vertex_location=first.get("vertex_location"), + ) + return total_cost + def vertex_ai_live_passthrough_handler( self, - websocket_messages: list[dict], - logging_obj, + websocket_messages: Sequence[object], + logging_obj: LiteLLMLoggingObj, url_route: str, start_time: datetime, end_time: datetime, @@ -274,28 +369,25 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): """ try: # Extract model from request body or kwargs - model: Final = kwargs.get("model", "gemini-2.0-flash-live-preview-04-09") + requested_model: Final = kwargs.get("model") + model: Final = ( + requested_model if isinstance(requested_model, str) else "gemini-2.0-flash-live-preview-04-09" + ) custom_llm_provider: Final = kwargs.get("custom_llm_provider", "vertex_ai") verbose_proxy_logger.debug( "Vertex AI Live API model: %s, custom_llm_provider: %s", model, custom_llm_provider ) - # Extract usage metadata from WebSocket messages - usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages) + usage: Final = self._session_usage(websocket_messages, model) - if not usage_metadata: + if usage is None: verbose_proxy_logger.warning("No usage metadata found in Vertex AI Live API WebSocket messages") return { "result": None, "kwargs": kwargs, } - # Create Usage object for standard LiteLLM logging - usage: Final = self._create_usage_object_from_metadata( - usage_metadata=usage_metadata, - grounding_metadata=_grounding_metadata(websocket_messages), - model=model, - ) + response_cost: Final = self._session_cost(websocket_messages, model, logging_obj) # Create a mock ModelResponse for standard logging litellm_model_response: Final = ModelResponse( @@ -306,6 +398,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): usage=usage, choices=[], ) + if response_cost is not None: + litellm_model_response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads the cost off the response's hidden params; the constructor's hidden_params kwarg is reset by pydantic kwargs["model"] = model kwargs["custom_llm_provider"] = custom_llm_provider @@ -314,7 +408,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): import re allowed_pattern: Final = re.compile(r"^[A-Za-z0-9._\-:]+$") - safe_model: Final = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]" + safe_model: Final = model if allowed_pattern.match(model) else "[REDACTED]" verbose_proxy_logger.debug( "Vertex AI Live API passthrough cost tracking - Model: %s, " "Prompt tokens: %s %s, Completion tokens: %s %s", diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index 1815ff134aa..dbb6c4cd702 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -24,7 +24,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.types.utils import LlmProviders, Usage +from litellm.types.utils import CostBreakdown, LlmProviders, Usage from litellm.proxy._types import UserAPIKeyAuth @@ -47,6 +47,7 @@ class TestVertexAILivePassthroughLoggingHandler: """Create a mock logging object""" mock = MagicMock(spec=LiteLLMLoggingObj) mock.model_call_details = {} + mock._response_cost_calculator.return_value = None return mock @pytest.fixture @@ -474,6 +475,64 @@ class TestVertexAILivePassthroughLoggingHandler: assert grounded > plain, "a grounded session must cost more than the same tokens ungrounded" + def _priced_logging_obj(self) -> LiteLLMLoggingObj: + """A real logging object, since the session's price is handed to it turn by turn.""" + logging_obj = LiteLLMLoggingObj( + model=self.NATIVE_AUDIO_MODEL, + messages=[], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="live-session", + function_id="live", + ) + logging_obj.update_environment_variables( + model=self.NATIVE_AUDIO_MODEL, + user="u", + optional_params={}, + litellm_params={}, + call_type="pass_through_endpoint", + ) + logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" + return logging_obj + + def _billed_session( + self, handler: VertexAILivePassthroughLoggingHandler, messages: list[dict[str, object]] + ) -> tuple[float, CostBreakdown]: + logging_obj = self._priced_logging_obj() + result = handler.vertex_ai_live_passthrough_handler( + websocket_messages=messages, + logging_obj=logging_obj, + url_route="/vertex_ai/live", + start_time=datetime.now(), + end_time=datetime.now(), + request_body={}, + model=self.NATIVE_AUDIO_MODEL, + custom_llm_provider="vertex_ai", + ) + assert result["result"] is not None, "the handler must produce a usage-bearing response to bill" + assert logging_obj.cost_breakdown is not None, "the session's price must reach the logging object" + return result["result"]._hidden_params["response_cost"], logging_obj.cost_breakdown + + def test_each_grounded_turn_pays_its_own_query_fee(self, handler): + """Google charges the grounding fee per grounded prompt, not per session. + + Summing the session into one usage collapsed two grounded turns into one query, so the + second question was answered for free. The bill now grows by one fee per grounded turn. + """ + head, turn = self._live_messages(self.AUDIO_SESSION[:1]) + grounding = self._grounding_frame({"webSearchQueries": ["q"]}) + + plain_cost, _ = self._billed_session(handler, [head, turn, turn]) + one_cost, one_breakdown = self._billed_session(handler, [head, grounding, turn, turn]) + two_cost, two_breakdown = self._billed_session(handler, [head, grounding, turn, grounding, turn]) + + fee = one_cost - plain_cost + assert fee > 0, "a grounded turn must cost more than the same tokens ungrounded" + assert two_cost - plain_cost == pytest.approx(2 * fee), "two grounded turns must pay the fee twice" + assert two_breakdown["total_cost"] == pytest.approx(two_cost) + assert two_breakdown["tool_usage_cost"] == pytest.approx(2 * one_breakdown["tool_usage_cost"]) + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): """Deliberate boundary: these tokens are reported here, and priced nowhere. @@ -676,6 +735,7 @@ class TestVertexAILivePassthroughIntegration: """Create a mock logging object""" mock = MagicMock(spec=LiteLLMLoggingObj) mock.model_call_details = {} + mock._response_cost_calculator.return_value = None return mock @patch( @@ -809,6 +869,7 @@ class TestVertexAILivePassthroughErrorHandling: """Create a mock logging object""" mock = MagicMock(spec=LiteLLMLoggingObj) mock.model_call_details = {} + mock._response_cost_calculator.return_value = None return mock def test_invalid_websocket_messages_format(self): From 710d4ae2a3aea57f78decbe856dca64e4b7e8224 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:33:31 -0700 Subject: [PATCH 024/107] fix(vertex-live): charge the fixed cost margin once per Live session --- ...tex_ai_live_passthrough_logging_handler.py | 18 +++++++++++---- .../test_vertex_ai_live_passthrough.py | 23 +++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 0ac04654e30..57ec960b032 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -314,14 +314,24 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): model: str, logging_obj: LiteLLMLoggingObj, ) -> float | None: - """Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice.""" + """Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice. + + The fixed cost margin is a flat per-request fee, so the session's single spend row carries it once + rather than once per turn. + """ turn_costs: Final = tuple(self._turn_cost(turn, model, logging_obj) for turn in _turns(websocket_messages)) priced: Final = tuple(turn_cost for turn_cost in turn_costs if turn_cost is not None) if not priced or len(priced) != len(turn_costs): return None breakdowns: Final = tuple(breakdown for _, breakdown in priced) first: Final = breakdowns[0] - total_cost: Final = sum(cost for cost, _ in priced) + fixed_margin: Final = first.get("margin_fixed_amount") or 0.0 + duplicated_fixed_margin: Final = fixed_margin * (len(priced) - 1) + total_cost: Final = sum(cost for cost, _ in priced) - duplicated_fixed_margin + summed_margin_total: Final = _summed(breakdowns, "margin_total_amount") + margin_total_amount: Final = ( + None if summed_margin_total is None else summed_margin_total - duplicated_fixed_margin + ) logging_obj.set_cost_breakdown( input_cost=_summed(breakdowns, "input_cost") or 0.0, output_cost=_summed(breakdowns, "output_cost") or 0.0, @@ -331,8 +341,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): discount_percent=first.get("discount_percent"), discount_amount=_summed(breakdowns, "discount_amount"), margin_percent=first.get("margin_percent"), - margin_fixed_amount=_summed(breakdowns, "margin_fixed_amount"), - margin_total_amount=_summed(breakdowns, "margin_total_amount"), + margin_fixed_amount=first.get("margin_fixed_amount"), + margin_total_amount=margin_total_amount, cache_read_cost=_summed(breakdowns, "cache_read_cost"), cache_creation_cost=_summed(breakdowns, "cache_creation_cost"), reasoning_cost=_summed(breakdowns, "reasoning_cost"), diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index dbb6c4cd702..70c2fda369b 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -13,6 +13,7 @@ from typing import Dict, List, Any, Optional import pytest import httpx +import litellm from typing_extensions import NotRequired, ReadOnly, TypedDict # Add the parent directory to the system path @@ -533,6 +534,28 @@ class TestVertexAILivePassthroughLoggingHandler: assert two_breakdown["total_cost"] == pytest.approx(two_cost) assert two_breakdown["tool_usage_cost"] == pytest.approx(2 * one_breakdown["tool_usage_cost"]) + def test_the_fixed_cost_margin_is_charged_once_per_session(self, handler): + """A fixed cost margin is a flat per-request fee, and a Live session is one spend row. + + Pricing each turn on its own applied the fixed margin per turn, so a two-turn session paid it + twice. The session now carries the fixed margin once no matter how many turns it billed. + """ + head, turn = self._live_messages(self.AUDIO_SESSION[:1]) + grounding = self._grounding_frame({"webSearchQueries": ["q"]}) + messages = [head, grounding, turn, grounding, turn] + + plain_cost, _ = self._billed_session(handler, messages) + + fixed_amount = 0.01 + with patch.object(litellm, "cost_margin_config", {"vertex_ai": {"fixed_amount": fixed_amount}}): + margined_cost, breakdown = self._billed_session(handler, messages) + + assert margined_cost - plain_cost == pytest.approx( + fixed_amount + ), "a two-turn session must add the fixed margin once, not once per billed turn" + assert breakdown["margin_fixed_amount"] == pytest.approx(fixed_amount) + assert breakdown["margin_total_amount"] == pytest.approx(fixed_amount) + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): """Deliberate boundary: these tokens are reported here, and priced nowhere. From fc62df33c2a0b049e8b48baad9bf97cd87aa57d3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 16:42:20 -0700 Subject: [PATCH 025/107] fix(proxy): resolve rotation and permission fields before the key policy and pin the effective-row contract --- .../key_management_endpoints.py | 38 +++++++---- .../key_management_endpoints.py | 13 +++- .../test_key_management_endpoints.py | 68 ++++++++++++++++++- 3 files changed, 100 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a1ce1700f68..841c22da050 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -359,7 +359,9 @@ def _effective_key_after_update( overlay: Final = MappingProxyType( {column: _decode_json_string_column(column, value) for column, value in non_default_values.items()} ) - return _verification_token_from_row(MappingProxyType({**existing_key_row.model_dump(), **overlay})) + return _verification_token_from_row( + MappingProxyType({**existing_key_row.model_dump(), **overlay, "object_permission": None}) + ) def _update_policy_request( @@ -381,7 +383,7 @@ def _update_policy_request( def _generate_budget_windows( budget_limits: Sequence[BudgetLimitEntry] | None, ) -> tuple[Mapping[str, object], ...] | None: - if budget_limits is None: + if not budget_limits: return None return tuple( MappingProxyType( @@ -402,15 +404,20 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite column_fields: Final = MappingProxyType( {field: value for field, value in requested.items() if field not in _KEY_METADATA_REQUEST_FIELDS} ) - request_metadata: Final = data.metadata or MappingProxyType({}) - folded_metadata: Final = {**request_metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict - columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates data_json in place + metadata: Final = data.metadata or MappingProxyType({}) + folded_metadata: Final = {**metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict + columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates in place expires: Final = ( now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None ) budget_reset_at: Final = ( get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None ) + key_rotation_at: Final = ( + now + timedelta(seconds=duration_in_seconds(duration=data.rotation_interval)) + if data.auto_rotate and data.rotation_interval + else None + ) return _verification_token_from_row( MappingProxyType( { @@ -418,6 +425,7 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite "metadata": encrypt_callback_vars(folded_metadata), "expires": expires, "budget_reset_at": budget_reset_at, + "key_rotation_at": key_rotation_at, "budget_limits": _generate_budget_windows(data.budget_limits), "object_permission": None, } @@ -3249,16 +3257,6 @@ async def update_key_fn( _enforce_upperbound_key_params(data, fill_defaults=False) non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row) - await _enforce_custom_key_policy( - hook=_custom_key_policy_hook(proxy_server), - build_policy_request=lambda: _update_policy_request( - operation="update", - existing_key_row=existing_key_row, - non_default_values=non_default_values, - request=data, - ), - ) - # Only validate key_alias format if it's actually being changed new_key_alias: Final = non_default_values.get("key_alias", None) if new_key_alias != existing_key_row.key_alias: @@ -3278,6 +3276,16 @@ async def update_key_fn( existing_key_alias=existing_key_row.key_alias, ) + await _enforce_custom_key_policy( + hook=_custom_key_policy_hook(proxy_server), + build_policy_request=lambda: _update_policy_request( + operation="update", + existing_key_row=existing_key_row, + non_default_values=non_default_values, + request=data, + ), + ) + if prisma_client is None: raise Exception("Not connected to DB!") diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 35abb8fb718..63bbaa5ba4e 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, model_validator from typing_extensions import ReadOnly, TypedDict @@ -128,11 +128,18 @@ class BulkUpdateTeamKeysRequest(BaseModel): return self -CustomKeyPolicyOperation = Literal["generate", "update", "regenerate"] +CustomKeyPolicyOperation: TypeAlias = Literal["generate", "update", "regenerate"] class CustomKeyPolicyRequest(LiteLLMPydanticObjectBase): - """What `general_settings.custom_key_policy` receives: the operation, the key row as it will be written, and the raw request.""" + """What `general_settings.custom_key_policy` receives. + + `effective_key` is the verification token row as it will be written: the existing row overlaid with the + requested changes, with `duration` resolved to `expires` and `budget_duration` to `budget_reset_at`. Values the + proxy fills in after the policy stay at their defaults: `token`, `key_name`, `created_by`, `updated_by` and the + soft-budget `budget_id` on generate, the rotated token on regenerate, and the `object_permission` relation on + every operation (`object_permission_id` is set; read `request.object_permission` for the requested change). + """ model_config = ConfigDict(protected_namespaces=(), frozen=True) 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 2b8c62523b9..b014e4428fd 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 @@ -31,6 +31,7 @@ from litellm.proxy._types import ( ResetSpendRequest, UpdateKeyRequest, ) +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -12362,7 +12363,9 @@ async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeyp mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) received: list[CustomKeyPolicyRequest] = [] policy = _seven_day_policy(received) - data = UpdateKeyRequest(key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0) + data = UpdateKeyRequest( + key=_POLICY_HASHED_TOKEN, duration="5d", max_budget=50.0, auto_rotate=True, rotation_interval="30d" + ) with ( patch( # test-quality-ok: cache eviction is outside the policy path @@ -12381,6 +12384,9 @@ async def test_update_key_fn_runs_custom_key_policy_on_the_effective_row(monkeyp mock_prisma_client.update_data.assert_awaited_once() assert len(received) == 1 _assert_update_policy_request(received[0], data) + key_rotation_at = received[0].effective_key.key_rotation_at + assert key_rotation_at is not None + assert abs(key_rotation_at - (datetime.now(timezone.utc) + timedelta(days=30))) < timedelta(seconds=60) @pytest.mark.asyncio @@ -12695,6 +12701,23 @@ async def test_effective_key_after_update_clears_expiry_for_a_minus_one_duration assert effective_key.expires is None +def test_effective_key_after_update_swaps_the_object_permission_id_and_drops_the_stale_relation(): + existing_key = LiteLLM_VerificationToken( + token="tok", + object_permission_id="op-old", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-old", mcp_servers=["old"]), + ) + + effective_key = _effective_key_after_update( + existing_key_row=existing_key, non_default_values={"object_permission_id": "op-new"} + ) + + assert effective_key.object_permission_id == "op-new" + assert effective_key.object_permission is None + assert existing_key.object_permission is not None + assert existing_key.object_permission.mcp_servers == ["old"] + + def test_effective_key_for_generate_reflects_the_processed_request_without_mutating_it(): now = datetime(2026, 1, 1, tzinfo=timezone.utc) data = GenerateKeyRequest( @@ -12705,12 +12728,21 @@ def test_effective_key_for_generate_reflects_the_processed_request_without_mutat tags=["t1"], budget_duration="1d", max_budget=3.0, + budget_limits=[{"budget_duration": "1d", "max_budget": 5.0}], + auto_rotate=True, + rotation_interval="30d", + object_permission={"mcp_servers": ["srv"]}, key_type=LiteLLMKeyType.LLM_API, ) effective_key = _effective_key_for_generate(data=data, now=now) assert effective_key.expires == now + timedelta(days=5) + assert effective_key.key_rotation_at == now + timedelta(days=30) + assert effective_key.budget_limits is not None + assert effective_key.budget_limits[0]["max_budget"] == 5.0 + assert effective_key.budget_limits[0]["reset_at"] is not None + assert effective_key.object_permission is None assert effective_key.org_id == "org-1" assert effective_key.metadata == {"a": 1, "guardrails": ["g1"], "tags": ["t1"]} assert effective_key.max_budget == 3.0 @@ -12722,6 +12754,18 @@ def test_effective_key_for_generate_reflects_the_processed_request_without_mutat assert data.guardrails == ["g1"] assert data.tags == ["t1"] assert data.duration == "5d" + assert data.budget_limits is not None + assert data.budget_limits[0].reset_at is None + assert data.object_permission is not None + assert data.object_permission.mcp_servers == ["srv"] + + +def test_effective_key_for_generate_stores_no_budget_windows_for_an_empty_list(): + effective_key = _effective_key_for_generate( + data=GenerateKeyRequest(budget_limits=[]), now=datetime(2026, 1, 1, tzinfo=timezone.utc) + ) + + assert effective_key.budget_limits is None def test_effective_key_for_generate_without_duration_never_expires(): @@ -12731,6 +12775,7 @@ def test_effective_key_for_generate_without_duration_never_expires(): assert effective_key.expires is None assert effective_key.budget_reset_at is None + assert effective_key.key_rotation_at is None assert effective_key.key_type == "default" @@ -12764,6 +12809,27 @@ async def test_enforce_custom_key_policy_uses_the_default_denial_message(): assert exc_info.value.detail == "Authentication Failed - Custom Auth Rule" +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_allows_when_the_decision_is_missing(): + received: list[CustomKeyPolicyRequest] = [] + + async def no_decision(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + received.append(policy_request) + return {} + + await _enforce_custom_key_policy(hook=no_decision, build_policy_request=_policy_request_for_generate) + + assert len(received) == 1 + assert received[0].operation == "generate" + + +@pytest.mark.asyncio +async def test_enforce_custom_key_policy_never_builds_the_request_without_a_hook(): + await _enforce_custom_key_policy( + hook=None, build_policy_request=lambda: pytest.fail("policy request built without a hook") + ) + + @pytest.mark.asyncio async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new_token(): """ From a159c7d98eb786a5bc7285dee0984512b6965a21 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 17:38:23 -0700 Subject: [PATCH 026/107] fix(proxy): write the requested object permission row only after the key policy allows the update --- .../key_management_endpoints.py | 38 ++-- .../test_key_management_endpoints.py | 164 +++++++++++++++++- 2 files changed, 178 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 841c22da050..ad2ab58d357 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2464,12 +2464,6 @@ async def prepare_key_update_data( # sentinel for Json? columns, so store the JSON literal null non_default_values["budget_limits"] = json.dumps(None) - if "object_permission" in non_default_values: - non_default_values = await _handle_update_object_permission( - data_json=non_default_values, - existing_key_row=existing_key_row, - ) - _metadata: Final = existing_key_row.metadata or {} # validate model_max_budget @@ -2490,13 +2484,12 @@ async def prepare_key_update_data( async def _handle_update_object_permission( data_json: dict, existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient, ) -> dict: - """ - Handle the update of object permission. - """ - from litellm.proxy.proxy_server import prisma_client + """Persist the requested object permission row and swap it for its id, only after the key policy allowed the write.""" + if "object_permission" not in data_json: + return data_json - # Use the common helper to handle the object permission update object_permission_id: Final = await handle_update_object_permission_common( data_json=data_json, existing_object_permission_id=existing_key_row.object_permission_id, @@ -2758,7 +2751,12 @@ async def _process_single_key_update( detail={"error": "Database not connected"}, ) - _data: Final = {**non_default_values, "token": update_key_request.key} + update_values: Final = await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + ) + _data: Final = {**update_values, "token": update_key_request.key} response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict "Mapping[str, object] | None", await prisma_client.update_data(token=update_key_request.key, data=_data), @@ -3289,18 +3287,23 @@ async def update_key_fn( if prisma_client is None: raise Exception("Not connected to DB!") + update_values: Final = await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + ) changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name response: Final = ( await _update_key_row_with_soft_budget( prisma_client=prisma_client, key=key, data=data, - non_default_values=non_default_values, + non_default_values=update_values, existing_key_row=existing_key_row, changed_by=changed_by, ) if "soft_budget" in data.model_fields_set - else await prisma_client.update_data(token=key, data=MappingProxyType({**non_default_values, "token": key})) + else await prisma_client.update_data(token=key, data=MappingProxyType({**update_values, "token": key})) ) # Delete - key from cache, since it's been updated! @@ -5324,7 +5327,12 @@ async def _execute_virtual_key_regeneration( request=data if data is not None else RegenerateKeyRequest(), ), ) - update_data.update(non_default_values) + update_values: Final = await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=key_in_db, + prisma_client=prisma_client, + ) + update_data.update(update_values) jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) # Snapshot before the token update: the FK cascade rewrites mapping rows to the new hash, 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 b014e4428fd..82458b5dfa3 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 @@ -1,3 +1,4 @@ +from collections.abc import Mapping from contextlib import ExitStack from typing import Final import json @@ -18,6 +19,7 @@ from litellm.proxy._types import ( GenerateKeyRequest, NewUserRequest, LiteLLM_BudgetTable, + LiteLLM_ObjectPermissionBase, LiteLLM_OrganizationTable, LiteLLM_ProjectTableCachedObj, LiteLLM_TeamTableCachedObj, @@ -1035,7 +1037,7 @@ async def test_key_generation_with_mcp_tool_permissions(monkeypatch): @pytest.mark.asyncio -async def test_key_update_object_permissions_existing_permission(monkeypatch): +async def test_key_update_object_permissions_existing_permission(): """ Test updating object permissions when a key already has an existing object_permission_id. @@ -1055,9 +1057,7 @@ async def test_key_update_object_permissions_existing_permission(monkeypatch): _handle_update_object_permission, ) - # Mock prisma client mock_prisma_client = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Mock existing key with object_permission_id existing_key_row = LiteLLM_VerificationToken( @@ -1097,6 +1097,7 @@ async def test_key_update_object_permissions_existing_permission(monkeypatch): result = await _handle_update_object_permission( data_json=data_json, existing_key_row=existing_key_row, + prisma_client=mock_prisma_client, ) # Verify the object_permission was removed from data_json and object_permission_id was set @@ -1111,7 +1112,7 @@ async def test_key_update_object_permissions_existing_permission(monkeypatch): @pytest.mark.asyncio -async def test_key_update_object_permissions_no_existing_permission(monkeypatch): +async def test_key_update_object_permissions_no_existing_permission(): """ Test creating object permissions when a key has no existing object_permission_id. @@ -1131,9 +1132,7 @@ async def test_key_update_object_permissions_no_existing_permission(monkeypatch) _handle_update_object_permission, ) - # Mock prisma client mock_prisma_client = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) existing_key_row_no_perm = LiteLLM_VerificationToken( token="test_token_hash_2", @@ -1164,6 +1163,7 @@ async def test_key_update_object_permissions_no_existing_permission(monkeypatch) result = await _handle_update_object_permission( data_json=data_json, existing_key_row=existing_key_row_no_perm, + prisma_client=mock_prisma_client, ) # Verify new object_permission_id was set @@ -1174,7 +1174,7 @@ async def test_key_update_object_permissions_no_existing_permission(monkeypatch) @pytest.mark.asyncio -async def test_key_update_object_permissions_missing_permission_record(monkeypatch): +async def test_key_update_object_permissions_missing_permission_record(): """ Test creating object permissions when existing object_permission_id record is not found. @@ -1194,9 +1194,7 @@ async def test_key_update_object_permissions_missing_permission_record(monkeypat _handle_update_object_permission, ) - # Mock prisma client mock_prisma_client = AsyncMock() - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) existing_key_row_missing_perm = LiteLLM_VerificationToken( token="test_token_hash_3", @@ -1227,6 +1225,7 @@ async def test_key_update_object_permissions_missing_permission_record(monkeypat result = await _handle_update_object_permission( data_json=data_json, existing_key_row=existing_key_row_missing_perm, + prisma_client=mock_prisma_client, ) # Verify new object_permission_id was set @@ -12470,6 +12469,153 @@ async def test_process_single_key_update_rejects_when_custom_key_policy_denies() assert [policy_request.operation for policy_request in received] == ["update"] +_OBJECT_PERMISSION_ID_AFTER_POLICY = "perm-after-policy" + + +def _record_object_permission_writes(mock_prisma_client: AsyncMock, events: list[str]) -> None: + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) + + async def upsert(**_kwargs: object) -> MagicMock: + events.append("permission row upsert") + return MagicMock(object_permission_id=_OBJECT_PERMISSION_ID_AFTER_POLICY) + + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(side_effect=upsert) + + +def _recording_policy(events: list[str], allowed: bool): + async def policy(policy_request: CustomKeyPolicyRequest) -> dict[str, object]: + events.append("policy") + return {"decision": allowed, "message": "key max_budget must be 1000 or less"} + + return policy + + +def _assert_permission_row_written_after_policy(events: list[str], written: Mapping[str, object]) -> None: + assert events == ["policy", "permission row upsert"] + assert written["object_permission_id"] == _OBJECT_PERMISSION_ID_AFTER_POLICY + assert "object_permission" not in written + + +def _assert_permission_row_untouched(mock_prisma_client: AsyncMock, events: list[str]) -> None: + assert events == ["policy"] + mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_awaited() + + +def _update_with_object_permission(max_budget: float) -> UpdateKeyRequest: + return UpdateKeyRequest( + key=_POLICY_HASHED_TOKEN, + max_budget=max_budget, + object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"]), + ) + + +def _setup_update_key_fn_object_permission_mocks(monkeypatch, allowed: bool) -> tuple[AsyncMock, list[str]]: + mock_prisma_client = _setup_update_key_fn_policy_mocks(monkeypatch, _policy_existing_team_key()) + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_policy", _recording_policy(events, allowed)) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", AsyncMock() + ) + return mock_prisma_client, events + + +async def _update_key_fn_with_object_permission(max_budget: float): + from litellm.proxy.management_endpoints.key_management_endpoints import update_key_fn + + return await update_key_fn( + request=MagicMock(), + data=_update_with_object_permission(max_budget=max_budget), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + ) + + +@pytest.mark.asyncio +async def test_update_key_fn_writes_the_object_permission_row_only_after_the_policy_allows(monkeypatch): + mock_prisma_client, events = _setup_update_key_fn_object_permission_mocks(monkeypatch, allowed=True) + + await _update_key_fn_with_object_permission(max_budget=50.0) + + _assert_permission_row_written_after_policy(events, mock_prisma_client.update_data.await_args.kwargs["data"]) + + +@pytest.mark.asyncio +async def test_update_key_fn_denied_by_the_policy_leaves_the_object_permission_row_untouched(monkeypatch): + mock_prisma_client, events = _setup_update_key_fn_object_permission_mocks(monkeypatch, allowed=False) + + with pytest.raises(ProxyException) as exc_info: + await _update_key_fn_with_object_permission(max_budget=5000.0) + + assert str(exc_info.value.code) == "403" + _assert_permission_row_untouched(mock_prisma_client, events) + mock_prisma_client.update_data.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_process_single_key_update_writes_the_object_permission_row_only_after_the_policy_allows(): + mock_prisma_client = AsyncMock() + updated_row = MagicMock() + updated_row.model_dump.return_value = {"max_budget": 50.0, "team_id": "team-a"} + mock_prisma_client.update_data = AsyncMock(return_value={"data": updated_row}) + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + + await _process_single_key_update_under_policy( + mock_prisma_client, _update_with_object_permission(max_budget=50.0), _recording_policy(events, allowed=True) + ) + + _assert_permission_row_written_after_policy(events, mock_prisma_client.update_data.await_args.kwargs["data"]) + + +@pytest.mark.asyncio +async def test_process_single_key_update_denied_by_the_policy_leaves_the_object_permission_row_untouched(): + mock_prisma_client = AsyncMock() + mock_prisma_client.update_data = AsyncMock() + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + + with pytest.raises(HTTPException) as exc_info: + await _process_single_key_update_under_policy( + mock_prisma_client, _update_with_object_permission(max_budget=5000.0), _recording_policy(events, allowed=False) + ) + + assert exc_info.value.status_code == 403 + _assert_permission_row_untouched(mock_prisma_client, events) + mock_prisma_client.update_data.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_regenerate_writes_the_object_permission_row_only_after_the_policy_allows(): + mock_prisma_client = _make_regenerate_mock_prisma() + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + data = RegenerateKeyRequest(max_budget=50.0, object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"])) + + with _regenerate_policy_mocks(_recording_policy(events, allowed=True), AsyncMock(), AsyncMock()): + await _regenerate_under_policy(mock_prisma_client, _make_regenerate_existing_key(), data) + + _assert_permission_row_written_after_policy( + events, mock_prisma_client.db.litellm_verificationtoken.update.await_args.kwargs["data"] + ) + + +@pytest.mark.asyncio +async def test_regenerate_denied_by_the_policy_leaves_the_object_permission_row_untouched(): + mock_prisma_client = _make_regenerate_mock_prisma() + events: list[str] = [] + _record_object_permission_writes(mock_prisma_client, events) + data = RegenerateKeyRequest(max_budget=5000.0, object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"])) + + with _regenerate_policy_mocks(_recording_policy(events, allowed=False), AsyncMock(), AsyncMock()): + with pytest.raises(HTTPException) as exc_info: + await _regenerate_under_policy(mock_prisma_client, _make_regenerate_existing_key(), data) + + assert exc_info.value.status_code == 403 + _assert_permission_row_untouched(mock_prisma_client, events) + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + + @pytest.mark.asyncio async def test_bulk_update_keys_runs_custom_key_policy_per_key(monkeypatch): from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys From 0429c645043e363ae2a87700a858349657850f6f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:22:31 -0700 Subject: [PATCH 027/107] fix(proxy): keep an empty duration out of the legacy update hook on key regenerate --- .../key_management_endpoints.py | 7 ++- .../test_key_management_endpoints.py | 57 ++++++++++++++++++- 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ad2ab58d357..8aaaba09e22 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -433,12 +433,17 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite ) +_EMPTY_DURATION_MEANS_UNCHANGED: Final = frozenset({"duration", "budget_duration"}) + + def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) -> UpdateKeyRequest | None: changed_fields: Final = MappingProxyType( { field: value for field, value in data.model_dump(exclude_unset=True).items() - if field in UpdateKeyRequest.model_fields and field != "key" + if field in UpdateKeyRequest.model_fields + and field != "key" + and not (field in _EMPTY_DURATION_MEANS_UNCHANGED and value == "") } ) if not changed_fields: 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 82458b5dfa3..df58424095f 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 @@ -12142,7 +12142,10 @@ async def test_execute_virtual_key_regeneration_allows_when_custom_key_update_ho @pytest.mark.asyncio -@pytest.mark.parametrize("data", [None, RegenerateKeyRequest()]) +@pytest.mark.parametrize( + "data", + [None, RegenerateKeyRequest(), RegenerateKeyRequest(duration=""), RegenerateKeyRequest(budget_duration="")], +) async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_without_changes(data): mock_prisma_client = _make_regenerate_mock_prisma() @@ -12184,6 +12187,58 @@ async def test_execute_virtual_key_regeneration_skips_custom_key_update_hook_wit assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_hides_the_untouched_modal_expiry_from_the_custom_key_update_hook(): + mock_prisma_client = _make_regenerate_mock_prisma() + untouched_modal_body = RegenerateKeyRequest( + key_alias=None, max_budget=None, tpm_limit=None, rpm_limit=None, duration="", grace_period="" + ) + received_data: list[UpdateKeyRequest] = [] + + async def hook(data: UpdateKeyRequest) -> dict[str, object]: + received_data.append(data) + if data.duration is not None and duration_in_seconds(data.duration) > duration_in_seconds("7d"): + return {"decision": False, "message": "duration must be <= 7d"} + return {"decision": True} + + with ( + patch( # test-quality-ok: deterministic token setup for the untouched modal body + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: grace-period path is outside the hook input + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: cache eviction is outside the hook input + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: rotation callback is outside the hook input + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_custom_key_update", hook), # test-quality-ok: inject policy hook + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=_make_regenerate_existing_key(), + hashed_api_key="abc123", + key="abc123", + data=untouched_modal_body, + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + assert len(received_data) == 1 + assert "duration" not in received_data[0].model_fields_set + assert received_data[0].model_fields_set >= {"key", "key_alias", "max_budget", "tpm_limit", "rpm_limit"} + + _POLICY_DENIAL_MESSAGE = "key duration must be 7d or less" _POLICY_HASHED_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" _POLICY_GENERATED_KEY = {"key": "sk-test-key", "expires": None, "user_id": "test-user", "team_id": None} From a635d7be6a10d5126790cd10df7b67fcfbf1a2a7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:29:09 -0700 Subject: [PATCH 028/107] fix(guardrails): write per-message guardrail rewrites back onto Responses input items A guardrail that answers one rewritten text per message it saw no longer matches the texts the Responses handler extracted once the request carries instructions or tool items, so the rewrite was rejected with a 500. Spread such an answer over the structured messages' text slots and write it back through the structured path, have Prompt Security modify return structured_messages directly, and give the chat completions pairing the same named rejection instead of a silent misalignment when the counts differ. --- .../base_llm/guardrail_translation/utils.py | 77 ++++++++++++- .../chat/guardrail_translation/handler.py | 4 + .../guardrail_translation/handler.py | 27 ++++- .../prompt_security/prompt_security.py | 39 ++++++- .../test_openai_guardrail_handler.py | 43 ++++++++ ...test_openai_responses_guardrail_handler.py | 103 ++++++++++++++++++ .../test_prompt_security_guardrails.py | 89 +++++++++++++++ 7 files changed, 373 insertions(+), 9 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 94a780f8148..383e668e45c 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -1,8 +1,10 @@ from __future__ import annotations import json -from collections.abc import Callable, Iterator, Sequence -from typing import Final, TypeVar +from collections.abc import Callable, Iterator, Mapping, Sequence +from itertools import accumulate +from types import MappingProxyType +from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles from pydantic import BaseModel @@ -364,3 +366,74 @@ def merge_guardrailed_scoped_messages( yield from appended return list(_merged()) + + +def _content_part_text(part: object) -> str | None: + if not isinstance(part, Mapping): + return None + text: Final = part.get("text") + return text if isinstance(text, str) else None + + +def message_text_slot_count(message: AllMessageValues) -> int: + content: Final = message.get("content") + if isinstance(content, str): + return 1 + if isinstance(content, list): + return sum(1 for part in content if _content_part_text(part) is not None) + return 0 + + +def _part_with_text(part: object, text: str) -> object: + if not isinstance(part, Mapping): + return part + return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts + + +def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> list[object]: + text_part_indices: Final = tuple( + index for index, part in enumerate(content) if _content_part_text(part) is not None + ) + replacement_by_index: Final = MappingProxyType(dict(zip(text_part_indices, texts))) + return [ # mutable-ok: message content stays a JSON list + _part_with_text(part, replacement_by_index[index]) if index in replacement_by_index else part + for index, part in enumerate(content) + ] + + +def _message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues: + content: Final = message.get("content") + if not isinstance(content, (str, list)): + return message + rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts) + rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts + return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped + + +def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None: + if message_text_slot_count(message) != len(texts): + return None + return _message_with_slot_texts(message, texts) + + +def messages_with_slot_texts( + messages: Sequence[AllMessageValues], + texts: Sequence[str], +) -> list[AllMessageValues] | None: + """Spread one flat list of rewritten texts over the messages' text slots, in order. + + A slot is a string ``content`` or one list part carrying a string ``text``; + images and other parts ride along untouched. A guardrail that answers one + text per message it saw produces exactly this shape, which stops matching + the endpoint's own per-text extraction as soon as the request carries + instructions or tool items. Returns None unless the counts line up exactly, + so a rewrite never lands on the wrong slot. + """ + slot_counts: Final = tuple(message_text_slot_count(message) for message in messages) + if sum(slot_counts) != len(texts): + return None + offsets: Final = tuple(accumulate(slot_counts, initial=0)) + return [ # mutable-ok: guardrail rows travel as a list + _message_with_slot_texts(message, texts[start:end]) + for message, start, end in zip(messages, offsets, offsets[1:]) + ] diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 58ff03e6a0d..56fda636e9a 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -196,6 +196,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): else: # Step 3: Map guardrail responses back to original message structure if guardrailed_texts and texts_to_check: + if len(guardrailed_texts) != len(text_task_mappings): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") await self._apply_guardrail_responses_to_input_texts( messages=messages, responses=guardrailed_texts, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 2fe11d9f7bd..61903a54cd3 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -53,6 +53,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( ) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, + messages_with_slot_texts, stream_item_field, stream_item_fingerprint, stream_item_items, @@ -395,6 +396,20 @@ def _patched_request_fields( ) +def _guardrailed_structured_messages( + structured_messages: Sequence[AllMessageValues] | None, + sent_text_count: int, + guardrailed_inputs: GenericGuardrailAPIInputs, +) -> Sequence[AllMessageValues] | None: + returned: Final = guardrailed_inputs.get("structured_messages") + if returned is not None and returned is not structured_messages: + return returned + rewritten_texts: Final = guardrailed_inputs.get("texts") + if not structured_messages or rewritten_texts is None or len(rewritten_texts) == sent_text_count: + return None + return messages_with_slot_texts(structured_messages, rewritten_texts) + + def _patch_or_convert_request_fields( raw_input: object, instructions: object, @@ -473,7 +488,8 @@ class OpenAIResponsesHandler(BaseTranslation): form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) ) extracted: Final = self._extract_guardrail_inputs(data, input_data, flattened_tool_groups) - if not extracted.inputs.get("texts"): + sent_texts: Final = extracted.inputs.get("texts") + if not sent_texts: return data if structured_messages: extracted.inputs["structured_messages"] = structured_messages @@ -486,7 +502,9 @@ class OpenAIResponsesHandler(BaseTranslation): 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) + written_back: Final = self._written_back_request_fields( + data, structured_messages, len(sent_texts), guardrailed_inputs + ) if written_back is not None: data["input"] = list(written_back.input) # mutable-ok: JSON body if written_back.instructions is None: @@ -553,10 +571,11 @@ class OpenAIResponsesHandler(BaseTranslation): def _written_back_request_fields( data: Mapping[str, object], structured_messages: Sequence[AllMessageValues] | None, + sent_text_count: int, guardrailed_inputs: GenericGuardrailAPIInputs, ) -> _RequestFields | None: - guardrailed: Final = guardrailed_inputs.get("structured_messages") - if guardrailed is None or guardrailed is structured_messages: + guardrailed: Final = _guardrailed_structured_messages(structured_messages, sent_text_count, guardrailed_inputs) + if guardrailed is None: return None return _patch_or_convert_request_fields( data.get("input"), diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 0954fe1698a..72c87a793a5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -2,6 +2,7 @@ import asyncio import base64 import os from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional import httpx @@ -14,11 +15,13 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.base_llm.guardrail_translation.utils import message_with_slot_texts from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -27,6 +30,7 @@ if TYPE_CHECKING: _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0 +_PROTECT_ROLES: Final = frozenset({"system", "user", "assistant"}) class PromptSecurityGuardrailMissingSecrets(Exception): @@ -275,14 +279,44 @@ class PromptSecurityGuardrail(CustomGuardrail): detail="Blocked by Prompt Security, Violations: " + ", ".join(violations), ) elif action == "modify": - # Extract modified texts from modified_messages modified_messages: Final = result.get("modified_messages", []) modified_texts: Final = self._extract_texts_from_messages(modified_messages) if modified_texts: inputs["texts"] = modified_texts + rewritten_messages: Final = self._structured_messages_with_modifications( + structured_messages, modified_messages + ) + if rewritten_messages is not None: + inputs["structured_messages"] = rewritten_messages return inputs + def _is_sent_to_protect(self, message: Mapping[str, object]) -> bool: + return self.check_tool_results or message.get("role") in _PROTECT_ROLES + + def _structured_messages_with_modifications( + self, + structured_messages: Sequence[AllMessageValues], + modified_messages: Sequence[Mapping[str, object]], + ) -> list[AllMessageValues] | None: + sent_indices: Final = tuple( + index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message) + ) + if not sent_indices or len(sent_indices) != len(modified_messages): + return None + rewritten: Final = tuple( + message_with_slot_texts(structured_messages[index], self._extract_texts_from_messages((modified,))) + for index, modified in zip(sent_indices, modified_messages) + ) + replacements: Final = MappingProxyType( + {index: message for index, message in zip(sent_indices, rewritten) if message is not None} + ) + if len(replacements) != len(sent_indices): + return None + return [ # mutable-ok: guardrail inputs take a list + replacements.get(index, message) for index, message in enumerate(structured_messages) + ] + async def _apply_guardrail_on_response( self, inputs: GenericGuardrailAPIInputs, @@ -678,14 +712,13 @@ class PromptSecurityGuardrail(CustomGuardrail): This allows checking tool results for indirect prompt injection when enabled. """ - supported_roles: Final = ["system", "user", "assistant"] filtered_messages: Final = [] transformed_count = 0 filtered_count = 0 for message in messages: role = message.get("role", "") - if role in supported_roles: + if role in _PROTECT_ROLES: filtered_messages.append(message) else: if self.check_tool_results: diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5a29a96829f..4e2291fdec1 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1893,6 +1893,49 @@ class TestScanOnlyToolResults: assert data["messages"][4]["content"] == "and then?" +class ToolDroppingTextGuardrail(CustomGuardrail): + """Answers one text per non-tool message it saw, the way a guardrail that + filters tool rows out before scanning does, and hands back only texts.""" + + def __init__(self): + super().__init__(guardrail_name="tool-dropping-redactor") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + kept = [m for m in inputs.get("structured_messages") or [] if m.get("role") != "tool"] + return {**inputs, "texts": [str(m.get("content")).replace("POISON", "[BLOCKED]") for m in kept]} + + +class TestPerMessageTextWriteBack: + """Texts that no longer pair one-to-one with what the handler extracted must be + rejected by name instead of sliding onto the wrong messages.""" + + @pytest.mark.asyncio + async def test_fewer_texts_than_extracted_over_a_tool_message_is_rejected(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + handler = OpenAIChatCompletionsHandler() + original_messages = [ + {"role": "system", "content": "SYSTEM-PROMPT"}, + {"role": "user", "content": "fetch the page"}, + {"role": "assistant", "content": "fetching"}, + {"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"}, + {"role": "user", "content": "and then?"}, + ] + data = {"messages": json.loads(json.dumps(original_messages))} + + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await handler.process_input_messages(data=data, guardrail_to_apply=ToolDroppingTextGuardrail()) + + assert excinfo.value.guardrail_name == "tool-dropping-redactor" + assert data["messages"] == original_messages, "a rejected rewrite must leave the request untouched" + + class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks""" 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 a4f0a77a9b6..b4b467773da 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 @@ -2338,6 +2338,109 @@ def _parallel_tool_call_input() -> list: ] +SSN = "123-45-6789" +REDACTED_SSN = "" + + +def _slot_texts(message: dict) -> list[str]: + content = message.get("content") + if isinstance(content, str): + return [content] + if isinstance(content, list): + return [part["text"] for part in content if isinstance(part, dict) and isinstance(part.get("text"), str)] + return [] + + +class PerMessageRedactionGuardrail(CustomGuardrail): + """Guardrail that answers one redacted text per message it was shown and hands + back only texts, the way Prompt Security in modify mode and a generic guardrail + API server that scans per message do.""" + + def __init__(self, extra_texts: int = 0): + super().__init__(guardrail_name="per-message-redactor") + self.extra_texts = extra_texts + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = inputs.get("structured_messages") or [] + texts = [text.replace(SSN, REDACTED_SSN) for message in messages for text in _slot_texts(message)] + return {**inputs, "texts": texts + ["junk"] * self.extra_texts} + + +class TestPerMessageTextWriteBack: + """A guardrail that rewrites one text per message it saw must land on the + instructions and the input items those messages came from, not be rejected.""" + + @pytest.mark.asyncio + async def test_instructions_plus_tool_replay_gets_each_rewrite_in_place(self): + handler = OpenAIResponsesHandler() + function_call_item = { + "type": "function_call", + "call_id": "call_1", + "name": "lookup_customer", + "arguments": '{"query": "' + SSN + '"}', + } + data = { + "model": "gpt-5.6", + "instructions": "Never repeat the SSN " + SSN + " back.", + "input": [ + {"role": "user", "content": "Look up " + SSN + " for me."}, + function_call_item, + {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, + ], + } + + result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) + + assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back." + assert [item.get("type", item.get("role")) for item in result["input"]] == [ + "user", + "function_call", + "function_call_output", + ] + assert _slot_texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."] + assert result["input"][1] == function_call_item + assert result["input"][2]["output"] == '{"ssn": "' + REDACTED_SSN + '"}' + assert result["input"][2]["call_id"] == "call_1" + + @pytest.mark.asyncio + async def test_string_input_with_instructions_keeps_the_two_apart(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "instructions": "Redact " + SSN + " everywhere.", + "input": "My SSN is " + SSN + ".", + } + + result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) + + assert result["instructions"] == "Redact " + REDACTED_SSN + " everywhere." + assert [_slot_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]] + + @pytest.mark.asyncio + async def test_count_matching_neither_texts_nor_messages_is_still_rejected(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + handler = OpenAIResponsesHandler() + original_input = [ + {"role": "user", "content": "Look up " + SSN + " for me."}, + {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, + ] + data = {"model": "gpt-5.6", "instructions": "Be terse.", "input": copy.deepcopy(original_input)} + + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await handler.process_input_messages(data, PerMessageRedactionGuardrail(extra_texts=1)) + + assert excinfo.value.guardrail_name == "per-message-redactor" + assert data["input"] == original_input + assert data["instructions"] == "Be terse." + + 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.""" diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index e650f796f29..9e83098eb04 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -174,6 +174,95 @@ async def test_apply_guardrail_modify_request(monkeypatch: pytest.MonkeyPatch): assert result["texts"] == ["User prompt with PII: SSN [REDACTED]"] +def _modify_response(modified_messages: list) -> Response: + mock_response = Response( + json={"result": {"prompt": {"action": "modify", "modified_messages": modified_messages}}}, + status_code=200, + request=Request(method="POST", url="https://test.prompt.security/api/protect"), + ) + mock_response.raise_for_status = lambda: None + return mock_response + + +def _tool_replay_messages() -> list: + return [ + {"role": "system", "content": "Never echo an SSN like 123-45-6789."}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Look up 123-45-6789"}, + {"type": "image_url", "image_url": {"url": "https://example.com/id-card.png"}}, + ], + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "123-45-6789"}'}, + {"role": "user", "content": "Summarize what you found."}, + ] + + +@pytest.mark.asyncio +async def test_modify_returns_structured_messages_with_tool_rows_kept(monkeypatch: pytest.MonkeyPatch): + """A per-message modify verdict comes back as structured_messages so the + endpoint handler can write it back by message, with the rows Prompt Security + never saw (tool results) and the non-text parts (images) left in place.""" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True) + messages = _tool_replay_messages() + inputs = {"texts": ["Look up 123-45-6789", "Summarize what you found."], "structured_messages": messages} + modified_messages = [ + {"role": "system", "content": "Never echo an SSN like [REDACTED]."}, + {"role": "user", "content": [{"type": "text", "text": "Look up [REDACTED]"}]}, + {"role": "assistant", "content": None}, + {"role": "user", "content": "Summarize what you found."}, + ] + + with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + + assert result["structured_messages"] == [ + {"role": "system", "content": "Never echo an SSN like [REDACTED]."}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Look up [REDACTED]"}, + {"type": "image_url", "image_url": {"url": "https://example.com/id-card.png"}}, + ], + }, + messages[2], + messages[3], + {"role": "user", "content": "Summarize what you found."}, + ] + assert result["structured_messages"] is not messages + assert result["texts"] == [ + "Never echo an SSN like [REDACTED].", + "Look up [REDACTED]", + "Summarize what you found.", + ] + + +@pytest.mark.asyncio +async def test_modify_with_unexpected_message_count_keeps_texts_only(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True) + messages = _tool_replay_messages() + inputs = {"texts": ["Look up 123-45-6789", "Summarize what you found."], "structured_messages": messages} + modified_messages = [{"role": "user", "content": "Look up [REDACTED]"}] + + with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): + result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + + assert result["structured_messages"] is messages + assert result["texts"] == ["Look up [REDACTED]"] + + @pytest.mark.asyncio async def test_apply_guardrail_allow_request(monkeypatch: pytest.MonkeyPatch): """Test that apply_guardrail allows safe prompts""" From 09314f239c34f35a86ce4155b8c4d445944a6038 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 00:54:29 -0700 Subject: [PATCH 029/107] fix(guardrails): hand per-message rewrites back as structured_messages A guardrail that rewrites text per chat message now returns the rewritten rows as structured_messages instead of only texts, so the Responses and chat handlers write the rewrite back through the structured path. The generic guardrail API response accepts an optional structured_messages list, Prompt Security modify builds one from modified_messages, and rows a server echoes back exactly as shown are restored to the original row objects because the request model drops undeclared keys. Texts-only per-message answers keep the named rejection on both endpoints. --- .../base_llm/guardrail_translation/utils.py | 42 ++--- .../guardrail_translation/handler.py | 27 +--- .../generic_guardrail_api.py | 30 +++- .../prompt_security/prompt_security.py | 8 +- .../guardrail_hooks/generic_guardrail_api.py | 15 +- ...test_openai_responses_guardrail_handler.py | 148 +++++++++--------- .../test_generic_guardrail_api.py | 105 +++++++++++++ .../test_prompt_security_guardrails.py | 14 +- 8 files changed, 244 insertions(+), 145 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 383e668e45c..1172c93959b 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -2,7 +2,6 @@ from __future__ import annotations import json from collections.abc import Callable, Iterator, Mapping, Sequence -from itertools import accumulate from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles @@ -390,7 +389,7 @@ def _part_with_text(part: object, text: str) -> object: return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts -def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> list[object]: +def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]: text_part_indices: Final = tuple( index for index, part in enumerate(content) if _content_part_text(part) is not None ) @@ -401,39 +400,18 @@ def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> ] -def _message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues: +def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None: + """Swap one rewritten text into each text slot of a chat row, in order. + + A slot is a string ``content`` or one list part carrying a string ``text``; + images and other parts ride along untouched. Returns None unless the counts + line up exactly, so a rewrite never lands on the wrong slot. + """ + if message_text_slot_count(message) != len(texts): + return None content: Final = message.get("content") if not isinstance(content, (str, list)): return message rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts) rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped - - -def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None: - if message_text_slot_count(message) != len(texts): - return None - return _message_with_slot_texts(message, texts) - - -def messages_with_slot_texts( - messages: Sequence[AllMessageValues], - texts: Sequence[str], -) -> list[AllMessageValues] | None: - """Spread one flat list of rewritten texts over the messages' text slots, in order. - - A slot is a string ``content`` or one list part carrying a string ``text``; - images and other parts ride along untouched. A guardrail that answers one - text per message it saw produces exactly this shape, which stops matching - the endpoint's own per-text extraction as soon as the request carries - instructions or tool items. Returns None unless the counts line up exactly, - so a rewrite never lands on the wrong slot. - """ - slot_counts: Final = tuple(message_text_slot_count(message) for message in messages) - if sum(slot_counts) != len(texts): - return None - offsets: Final = tuple(accumulate(slot_counts, initial=0)) - return [ # mutable-ok: guardrail rows travel as a list - _message_with_slot_texts(message, texts[start:end]) - for message, start, end in zip(messages, offsets, offsets[1:]) - ] diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 61903a54cd3..2fe11d9f7bd 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -53,7 +53,6 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( ) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, - messages_with_slot_texts, stream_item_field, stream_item_fingerprint, stream_item_items, @@ -396,20 +395,6 @@ def _patched_request_fields( ) -def _guardrailed_structured_messages( - structured_messages: Sequence[AllMessageValues] | None, - sent_text_count: int, - guardrailed_inputs: GenericGuardrailAPIInputs, -) -> Sequence[AllMessageValues] | None: - returned: Final = guardrailed_inputs.get("structured_messages") - if returned is not None and returned is not structured_messages: - return returned - rewritten_texts: Final = guardrailed_inputs.get("texts") - if not structured_messages or rewritten_texts is None or len(rewritten_texts) == sent_text_count: - return None - return messages_with_slot_texts(structured_messages, rewritten_texts) - - def _patch_or_convert_request_fields( raw_input: object, instructions: object, @@ -488,8 +473,7 @@ class OpenAIResponsesHandler(BaseTranslation): form.chat_tools for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(original_tools) ) extracted: Final = self._extract_guardrail_inputs(data, input_data, flattened_tool_groups) - sent_texts: Final = extracted.inputs.get("texts") - if not sent_texts: + if not extracted.inputs.get("texts"): return data if structured_messages: extracted.inputs["structured_messages"] = structured_messages @@ -502,9 +486,7 @@ class OpenAIResponsesHandler(BaseTranslation): 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, len(sent_texts), guardrailed_inputs - ) + 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: @@ -571,11 +553,10 @@ class OpenAIResponsesHandler(BaseTranslation): def _written_back_request_fields( data: Mapping[str, object], structured_messages: Sequence[AllMessageValues] | None, - sent_text_count: int, guardrailed_inputs: GenericGuardrailAPIInputs, ) -> _RequestFields | None: - guardrailed: Final = _guardrailed_structured_messages(structured_messages, sent_text_count, guardrailed_inputs) - if guardrailed is None: + guardrailed: Final = guardrailed_inputs.get("structured_messages") + if guardrailed is None or guardrailed is structured_messages: return None return _patch_or_convert_request_fields( data.get("input"), diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index d8296003ae9..16159d32a7f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -7,7 +7,7 @@ import fnmatch import os -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional import httpx @@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.llms.openai import ChatCompletionToolParam +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest, @@ -150,6 +150,22 @@ def _extract_inbound_headers( return None +def _rows_with_unchanged_originals( + original_rows: Sequence[AllMessageValues] | None, + shown_rows: Sequence[AllMessageValues] | None, + returned_rows: Sequence[AllMessageValues], +) -> tuple[AllMessageValues, ...]: + """The request model drops row keys its message types do not declare, so a + row the server echoes back verbatim is restored to the original row object; + only rows the server actually changed reach the endpoint write-back.""" + if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows): + return tuple(returned_rows) + return tuple( + original if returned == shown else returned + for original, shown, returned in zip(original_rows, shown_rows, returned_rows) + ) + + class GenericGuardrailAPI(CustomGuardrail): """ Generic Guardrail API integration for LiteLLM. @@ -322,6 +338,8 @@ class GenericGuardrailAPI(CustomGuardrail): texts: list, images: list[str] | None, tools: list[ChatCompletionToolParam] | None, + structured_messages: Sequence[AllMessageValues] | None, + shown_messages: Sequence[AllMessageValues] | None, guardrail_response: GenericGuardrailAPIResponse, ) -> GenericGuardrailAPIInputs: # Action is NONE or no modifications needed @@ -336,6 +354,12 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools + if guardrail_response.structured_messages: + return_inputs["structured_messages"] = list( # mutable-ok: guardrail inputs take a list + _rows_with_unchanged_originals( + structured_messages, shown_messages, guardrail_response.structured_messages + ) + ) if guardrail_response.stream_holdback_chars is not None: return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars return return_inputs @@ -473,6 +497,8 @@ class GenericGuardrailAPI(CustomGuardrail): texts=texts, images=images, tools=tools, + structured_messages=structured_messages, + shown_messages=guardrail_request.structured_messages, guardrail_response=guardrail_response, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 72c87a793a5..41d3b202344 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -287,7 +287,7 @@ class PromptSecurityGuardrail(CustomGuardrail): structured_messages, modified_messages ) if rewritten_messages is not None: - inputs["structured_messages"] = rewritten_messages + inputs["structured_messages"] = list(rewritten_messages) # mutable-ok: guardrail inputs take a list return inputs @@ -298,7 +298,7 @@ class PromptSecurityGuardrail(CustomGuardrail): self, structured_messages: Sequence[AllMessageValues], modified_messages: Sequence[Mapping[str, object]], - ) -> list[AllMessageValues] | None: + ) -> tuple[AllMessageValues, ...] | None: sent_indices: Final = tuple( index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message) ) @@ -313,9 +313,7 @@ class PromptSecurityGuardrail(CustomGuardrail): ) if len(replacements) != len(sent_indices): return None - return [ # mutable-ok: guardrail inputs take a list - replacements.get(index, message) for index, message in enumerate(structured_messages) - ] + return tuple(replacements.get(index, message) for index, message in enumerate(structured_messages)) async def _apply_guardrail_on_response( self, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 4a868c48352..44e2cc2404f 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -1,4 +1,5 @@ -from typing import Any, Final, Literal +from collections.abc import Mapping, Sequence +from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles from pydantic import BaseModel, ConfigDict, Field from typing_extensions import TypedDict @@ -158,12 +159,21 @@ def coerce_stream_holdback_value(value: Any) -> int: return 0 +def structured_messages_from_response(value: object) -> Sequence[AllMessageValues] | None: + if not isinstance(value, list): + return None + if not all(isinstance(message, Mapping) and isinstance(message.get("role"), str) for message in value): + return None + return cast("Sequence[AllMessageValues]", value) # cast-ok: JSON rows checked for a role, the same trust texts get + + class GenericGuardrailAPIResponse: """Response model for the Generic Guardrail API""" texts: list[str] | None images: list[str] | None tools: list[GuardrailToolParam] | None + structured_messages: Sequence[AllMessageValues] | None action: str blocked_reason: str | None stream_holdback_chars: list[int] | None @@ -176,12 +186,14 @@ class GenericGuardrailAPIResponse: images: list[str] | None = None, tools: list[GuardrailToolParam] | None = None, stream_holdback_chars: list[int] | None = None, + structured_messages: Sequence[AllMessageValues] | None = None, ) -> None: self.action = action self.blocked_reason = blocked_reason self.texts = texts self.images = images self.tools = tools + self.structured_messages = structured_messages # Number of trailing chars, indexed the same as ``texts``, that the # framework must withhold from streaming emission until the next # processing round (word-boundary safety for text transformations). @@ -200,4 +212,5 @@ class GenericGuardrailAPIResponse: images=data.get("images"), tools=data.get("tools"), stream_holdback_chars=stream_holdback_chars, + structured_messages=structured_messages_from_response(data.get("structured_messages")), ) 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 b4b467773da..394f134e99c 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 @@ -8,7 +8,7 @@ with guardrail transformations. import copy from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import logging @@ -31,6 +31,7 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, @@ -2342,103 +2343,94 @@ SSN = "123-45-6789" REDACTED_SSN = "" -def _slot_texts(message: dict) -> list[str]: - content = message.get("content") - if isinstance(content, str): - return [content] - if isinstance(content, list): - return [part["text"] for part in content if isinstance(part, dict) and isinstance(part.get("text"), str)] - return [] +def _redacted(value: object) -> object: + if isinstance(value, str): + return value.replace(SSN, REDACTED_SSN) + if isinstance(value, list): + return [{**part, "text": _redacted(part["text"])} if "text" in part else part for part in value] + return value -class PerMessageRedactionGuardrail(CustomGuardrail): - """Guardrail that answers one redacted text per message it was shown and hands - back only texts, the way Prompt Security in modify mode and a generic guardrail - API server that scans per message do.""" +def _per_message_guardrail_server(structured_messages_in_answer: bool) -> Callable[..., MagicMock]: + """Answers one redacted text per chat row it was shown, the way a guardrail + that scans per message does, and optionally the rewritten rows themselves.""" - def __init__(self, extra_texts: int = 0): - super().__init__(guardrail_name="per-message-redactor") - self.extra_texts = extra_texts + def post(url: str, json: dict, headers: dict) -> MagicMock: + rows = json["structured_messages"] + answer: dict = { + "action": "GUARDRAIL_INTERVENED", + "texts": [_redacted(row["content"]) if isinstance(row.get("content"), str) else "" for row in rows], + } + if structured_messages_in_answer: + answer["structured_messages"] = [{**row, "content": _redacted(row.get("content"))} for row in rows] + response = MagicMock() + response.json.return_value = answer + response.raise_for_status = MagicMock() + return response - async def apply_guardrail( - self, - inputs: GenericGuardrailAPIInputs, - request_data: dict, - input_type: Literal["request", "response"], - logging_obj: Optional[Any] = None, - ) -> GenericGuardrailAPIInputs: - messages = inputs.get("structured_messages") or [] - texts = [text.replace(SSN, REDACTED_SSN) for message in messages for text in _slot_texts(message)] - return {**inputs, "texts": texts + ["junk"] * self.extra_texts} + return post -class TestPerMessageTextWriteBack: - """A guardrail that rewrites one text per message it saw must land on the - instructions and the input items those messages came from, not be rejected.""" +def _per_message_redactor() -> GenericGuardrailAPI: + return GenericGuardrailAPI( + api_base="https://guardrail.test", + guardrail_name="per-message-redactor", + event_hook="pre_call", + default_on=True, + ) + + +def _tool_replay_request() -> dict: + return { + "model": "gpt-5.6", + "instructions": "Never repeat the SSN " + SSN + " back.", + "input": [ + {"role": "user", "content": "Look up " + SSN + " for me."}, + {"type": "function_call", "call_id": "call_1", "name": "lookup_customer", "arguments": '{"id": "42"}'}, + {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, + ], + } + + +class TestPerMessageRewriteWriteBack: + """A guardrail that rewrites per chat row hands the rows back as + structured_messages, and the handler lands them on the instructions and the + input items they came from; the same rewrite handed back as texts alone has + no item to land on and is rejected by name instead of sent unrewritten.""" @pytest.mark.asyncio - async def test_instructions_plus_tool_replay_gets_each_rewrite_in_place(self): - handler = OpenAIResponsesHandler() - function_call_item = { - "type": "function_call", - "call_id": "call_1", - "name": "lookup_customer", - "arguments": '{"query": "' + SSN + '"}', - } - data = { - "model": "gpt-5.6", - "instructions": "Never repeat the SSN " + SSN + " back.", - "input": [ - {"role": "user", "content": "Look up " + SSN + " for me."}, - function_call_item, - {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, - ], - } + async def test_structured_rows_land_on_instructions_and_tool_output(self): + guardrail = _per_message_redactor() + data = _tool_replay_request() + function_call_item = data["input"][1] - result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)): + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back." - assert [item.get("type", item.get("role")) for item in result["input"]] == [ - "user", - "function_call", - "function_call_output", - ] - assert _slot_texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."] + assert _texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."] assert result["input"][1] == function_call_item - assert result["input"][2]["output"] == '{"ssn": "' + REDACTED_SSN + '"}' - assert result["input"][2]["call_id"] == "call_1" - - @pytest.mark.asyncio - async def test_string_input_with_instructions_keeps_the_two_apart(self): - handler = OpenAIResponsesHandler() - data = { - "model": "gpt-5.6", - "instructions": "Redact " + SSN + " everywhere.", - "input": "My SSN is " + SSN + ".", + assert result["input"][2] == { + "type": "function_call_output", + "call_id": "call_1", + "output": '{"ssn": "' + REDACTED_SSN + '"}', } - result = await handler.process_input_messages(data, PerMessageRedactionGuardrail()) - - assert result["instructions"] == "Redact " + REDACTED_SSN + " everywhere." - assert [_slot_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]] - @pytest.mark.asyncio - async def test_count_matching_neither_texts_nor_messages_is_still_rejected(self): + async def test_texts_only_per_message_answer_is_rejected_by_name(self): from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite - handler = OpenAIResponsesHandler() - original_input = [ - {"role": "user", "content": "Look up " + SSN + " for me."}, - {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'}, - ] - data = {"model": "gpt-5.6", "instructions": "Be terse.", "input": copy.deepcopy(original_input)} + guardrail = _per_message_redactor() + data = _tool_replay_request() + original = copy.deepcopy(data) - with pytest.raises(UnappliableRequestRewrite) as excinfo: - await handler.process_input_messages(data, PerMessageRedactionGuardrail(extra_texts=1)) + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)): + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await OpenAIResponsesHandler().process_input_messages(data, guardrail) assert excinfo.value.guardrail_name == "per-message-redactor" - assert data["input"] == original_input - assert data["instructions"] == "Be terse." + assert data["input"] == original["input"] + assert data["instructions"] == original["instructions"] class TestProvenancePatching: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 83cc9ae8bb9..cc9942e0e40 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -582,6 +582,111 @@ class TestGuardrailActions: assert result_images is None +class TestStructuredMessagesInResponse: + """A guardrail server that rewrites per chat row answers with the rewritten + rows as structured_messages, which the endpoint handlers write back by row.""" + + @pytest.mark.asyncio + async def test_returned_rows_are_handed_back_as_structured_messages( + self, generic_guardrail, mock_request_data_input + ): + rewritten_rows = [ + {"role": "system", "content": "Never repeat an SSN."}, + {"role": "user", "content": "Look up [REDACTED] for me."}, + {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'}, + ] + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Never repeat an SSN.", "Look up [REDACTED] for me.", '{"ssn": "[REDACTED]"}'], + "structured_messages": rewritten_rows, + } + mock_response.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."]}, + request_data=mock_request_data_input, + input_type="request", + ) + + assert guardrailed_inputs["structured_messages"] == rewritten_rows + assert guardrailed_inputs["texts"] == mock_response.json.return_value["texts"] + + @pytest.mark.asyncio + async def test_rows_echoed_back_as_shown_keep_their_original_keys( + self, generic_guardrail, mock_request_data_input + ): + tool_call_row = { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}, "index": 0} + ], + } + original_rows = [ + {"role": "user", "content": "Look up 123-45-6789 for me.", "name": "pat"}, + tool_call_row, + {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "123-45-6789"}'}, + ] + + def echo_with_tool_output_redacted(url, json, headers): + shown_rows = json["structured_messages"] + assert "index" not in shown_rows[1]["tool_calls"][0] + assert "name" not in shown_rows[0] + answer = MagicMock() + answer.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Look up 123-45-6789 for me."], + "structured_messages": [ + shown_rows[0], + shown_rows[1], + {**shown_rows[2], "content": '{"ssn": "[REDACTED]"}'}, + ], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object(generic_guardrail.async_handler, "post", side_effect=echo_with_tool_output_redacted): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."], "structured_messages": original_rows}, + request_data=mock_request_data_input, + input_type="request", + ) + + returned_rows = guardrailed_inputs["structured_messages"] + assert returned_rows[0] is original_rows[0] + assert returned_rows[1] is tool_call_row + assert returned_rows[2] == {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "structured_messages", + [[], [{"content": "a row with no role"}], "not a list"], + ids=["empty", "no_role", "not_a_list"], + ) + async def test_rows_that_are_not_chat_messages_are_ignored( + self, generic_guardrail, mock_request_data_input, structured_messages + ): + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["[REDACTED]"], + "structured_messages": structured_messages, + } + mock_response.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Look up 123-45-6789 for me."]}, + request_data=mock_request_data_input, + input_type="request", + ) + + assert "structured_messages" not in guardrailed_inputs + assert guardrailed_inputs["texts"] == ["[REDACTED]"] + + class TestImageSupport: """Test image handling in guardrail requests""" diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index 9e83098eb04..70083e50f01 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -1,5 +1,6 @@ import asyncio import base64 +from collections.abc import Mapping, Sequence from unittest.mock import AsyncMock, patch import pytest @@ -12,6 +13,7 @@ from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security im PromptSecurityGuardrailMissingSecrets, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.types.llms.openai import AllMessageValues def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch): @@ -174,7 +176,7 @@ async def test_apply_guardrail_modify_request(monkeypatch: pytest.MonkeyPatch): assert result["texts"] == ["User prompt with PII: SSN [REDACTED]"] -def _modify_response(modified_messages: list) -> Response: +def _modify_response(modified_messages: Sequence[Mapping[str, object]]) -> Response: mock_response = Response( json={"result": {"prompt": {"action": "modify", "modified_messages": modified_messages}}}, status_code=200, @@ -184,7 +186,7 @@ def _modify_response(modified_messages: list) -> Response: return mock_response -def _tool_replay_messages() -> list: +def _tool_replay_messages() -> list[AllMessageValues]: return [ {"role": "system", "content": "Never echo an SSN like 123-45-6789."}, { @@ -224,7 +226,9 @@ async def test_modify_returns_structured_messages_with_tool_rows_kept(monkeypatc ] with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): - result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={"messages": messages}, input_type="request" + ) assert result["structured_messages"] == [ {"role": "system", "content": "Never echo an SSN like [REDACTED]."}, @@ -257,7 +261,9 @@ async def test_modify_with_unexpected_message_count_keeps_texts_only(monkeypatch modified_messages = [{"role": "user", "content": "Look up [REDACTED]"}] with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): - result = await guardrail.apply_guardrail(inputs=inputs, request_data={"messages": messages}, input_type="request") + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={"messages": messages}, input_type="request" + ) assert result["structured_messages"] is messages assert result["texts"] == ["Look up [REDACTED]"] From 93f3911b32e04c124027e8c5a12961fdcb0fca90 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:22:49 -0700 Subject: [PATCH 030/107] fix(guardrails): drop the types import CodeQL reads as a package cycle --- litellm/llms/base_llm/guardrail_translation/utils.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 1172c93959b..34d648cf184 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -2,7 +2,6 @@ from __future__ import annotations import json from collections.abc import Callable, Iterator, Mapping, Sequence -from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles from pydantic import BaseModel @@ -390,13 +389,10 @@ def _part_with_text(part: object, text: str) -> object: def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]: - text_part_indices: Final = tuple( - index for index, part in enumerate(content) if _content_part_text(part) is not None - ) - replacement_by_index: Final = MappingProxyType(dict(zip(text_part_indices, texts))) + remaining_texts: Final = iter(texts) return [ # mutable-ok: message content stays a JSON list - _part_with_text(part, replacement_by_index[index]) if index in replacement_by_index else part - for index, part in enumerate(content) + _part_with_text(part, next(remaining_texts)) if _content_part_text(part) is not None else part + for part in content ] From 1760fe628436fac7fb292fae9ecdf6016a51fc55 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:22:49 -0700 Subject: [PATCH 031/107] refactor(guardrails): return a fresh inputs mapping from the Prompt Security modify branch --- .../prompt_security/prompt_security.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 41d3b202344..5d533e98ff6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -37,6 +37,18 @@ class PromptSecurityGuardrailMissingSecrets(Exception): pass +def _inputs_with_modifications( + inputs: GenericGuardrailAPIInputs, + modified_texts: list[str], + rewritten_messages: Sequence[AllMessageValues] | None, +) -> GenericGuardrailAPIInputs: + texts_patch: Final[GenericGuardrailAPIInputs] = {"texts": modified_texts} if modified_texts else {} + messages_patch: Final[GenericGuardrailAPIInputs] = ( + {"structured_messages": list(rewritten_messages)} if rewritten_messages is not None else {} + ) + return {**inputs, **texts_patch, **messages_patch} + + class _ProtectVerdict(TypedDict, total=False): """One side (``prompt`` or ``response``) of an ``/api/protect`` verdict.""" @@ -280,14 +292,11 @@ class PromptSecurityGuardrail(CustomGuardrail): ) elif action == "modify": modified_messages: Final = result.get("modified_messages", []) - modified_texts: Final = self._extract_texts_from_messages(modified_messages) - if modified_texts: - inputs["texts"] = modified_texts - rewritten_messages: Final = self._structured_messages_with_modifications( - structured_messages, modified_messages + return _inputs_with_modifications( + inputs, + self._extract_texts_from_messages(modified_messages), + self._structured_messages_with_modifications(structured_messages, modified_messages), ) - if rewritten_messages is not None: - inputs["structured_messages"] = list(rewritten_messages) # mutable-ok: guardrail inputs take a list return inputs From 23a98cb85134342e20ccb7387875c9b0037bc7d2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:22:49 -0700 Subject: [PATCH 032/107] chore(ui): regenerate schema.d.ts after merging main --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From cccff657cf376843a1e9c732e9ec33c46183e01e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 01:35:19 -0700 Subject: [PATCH 033/107] fix(guardrails): scan the Anthropic top-level system prompt and tool_use arguments --- .../chat/guardrail_translation/handler.py | 224 +++++++++++++--- .../test_anthropic_guardrail_handler.py | 248 ++++++++++++++++-- 2 files changed, 413 insertions(+), 59 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50345d70d..eb16278c9bd 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -103,9 +103,24 @@ class ToolResultBlockTextTarget: block_idx: int -InputWriteBackTarget = ( - MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget -) +@dataclass(frozen=True, slots=True) +class SystemStringTarget: + pass + + +@dataclass(frozen=True, slots=True) +class SystemBlockTextTarget: + block_idx: int + + +@dataclass(frozen=True, slots=True) +class ToolUseInputTarget: + msg_idx: int + content_idx: int + + +MessageTextTarget = MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget +InputWriteBackTarget = SystemStringTarget | SystemBlockTextTarget | MessageTextTarget def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]: @@ -146,10 +161,17 @@ class ScannedText: target: InputWriteBackTarget +@dataclass(frozen=True, slots=True) +class ScannedToolCall: + tool_call: ChatCompletionToolCallChunk + target: ToolUseInputTarget + + @dataclass(frozen=True, slots=True) class ExtractedInput: scanned: tuple[ScannedText, ...] images: tuple[str, ...] + tool_calls: tuple[ScannedToolCall, ...] = () EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=()) @@ -161,6 +183,76 @@ class _ToolCallShape: arguments: str +def _is_client_tool_use(block: Mapping[str, object]) -> bool: + return ( + block.get("type") == "tool_use" + and isinstance(block.get("id"), str) + and isinstance(block.get("name"), str) + and isinstance(block.get("input"), Mapping) + ) + + +def _write_back_system_block(system: object, block_idx: int, response: str) -> None: + if not isinstance(system, list): + return + text_blocks: Final = tuple(block for block in system if isinstance(block, dict) and block.get("type") == "text") + if block_idx < len(text_blocks): + text_blocks[block_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + + +def _write_back_message_text(message: _WritableMessage, target: MessageTextTarget, response: str) -> None: + content: Final = message.get("content", None) + if content is None: + return + match target: + case MessageContentTarget(): + if isinstance(content, str): + message["content"] = response # mutable-ok: guardrails rewrite the caller's request payload in place + case ContentBlockTextTarget(content_idx=content_idx): + if isinstance(content, list): + content[content_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case ToolResultStringTarget(content_idx=content_idx): + if isinstance(content, list): + content[content_idx]["content"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx): + if isinstance(content, list): + content[content_idx]["content"][block_idx]["text"] = ( + response # mutable-ok: guardrails rewrite the caller's request payload in place + ) + case _: + assert_never(target) + + +def _write_back_tool_use(message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape) -> None: + content: Final = message.get("content", None) + block: Final = content[target.content_idx] if isinstance(content, list) else None + if not isinstance(block, dict): + return + try: + rewritten_input: Final = json.loads(shape.arguments) + except json.JSONDecodeError: + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail returned non-JSON arguments for tool_use %s; keeping its input", + block.get("id"), + ) + return + if not isinstance(rewritten_input, dict): + verbose_proxy_logger.warning( + "Anthropic Messages: guardrail returned non-object arguments for tool_use %s; keeping its input", + block.get("id"), + ) + return + block["input"] = rewritten_input # mutable-ok: guardrails rewrite the caller's request payload in place + if shape.name is not None and shape.name != block.get("name"): + block["name"] = shape.name # mutable-ok: guardrails rewrite the caller's request payload in place + + @dataclass(frozen=True, slots=True) class _SSEFieldRewrite: """One field of one nested section of a buffered SSE event, rewritten.""" @@ -452,9 +544,8 @@ class AnthropicMessagesHandler(BaseTranslation): skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply) scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply) - # Exclude only the trusted top-level prompt. In-sequence system entries are untrusted - # and must stay aligned with texts_to_check for positional masking. When the top-level - # prompt is included, the pre-existing count mismatch disables positional masking. + # The top-level prompt is translated on its own below so it can be hoisted in front of + # any mid-turn system entries and scanned first, aligned with that structured position. translation_source: Final = { # mutable-ok: API message payload key: value for key, value in data.items() if key != "system" } @@ -490,7 +581,12 @@ class AnthropicMessagesHandler(BaseTranslation): ] ) - # Step 1: Extract all text content and images + # Step 1: Extract all text content, images, and tool calls + top_level_system_scanned: Final = ( + () + if hoisted_system_message is None or scan_only_tool_results + else self._extract_top_level_system_text(hoisted_system_message) + ) extracted: Final = tuple( self._extract_input_text_and_images( message=message, @@ -501,17 +597,27 @@ class AnthropicMessagesHandler(BaseTranslation): ) for msg_idx, message in enumerate(messages) ) - scanned: Final = tuple(item for one_message in extracted for item in one_message.scanned) + scanned: Final = ( + *top_level_system_scanned, + *(item for one_message in extracted for item in one_message.scanned), + ) texts_to_check: Final = [item.text for item in scanned] # mutable-ok: GenericGuardrailAPIInputs takes list[str] images_to_check: Final = [ image for one_message in extracted for image in one_message.images ] # mutable-ok: GenericGuardrailAPIInputs takes list[str] + scanned_tool_calls: Final = tuple(item for one_message in extracted for item in one_message.tool_calls) + tool_calls_to_check: Final = [ + item.tool_call for item in scanned_tool_calls + ] # mutable-ok: GenericGuardrailAPIInputs takes list[ChatCompletionToolCallChunk] + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) - # Step 2: Apply guardrail to all texts in batch - if texts_to_check: + # Step 2: Apply guardrail to all texts and tool calls in batch + if texts_to_check or tool_calls_to_check: inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check + if tool_calls_to_check: + inputs["tool_calls"] = tool_calls_to_check if tools_to_check: inputs["tools"] = tools_to_check original_structured_messages: Final = structured_messages @@ -572,10 +678,16 @@ class AnthropicMessagesHandler(BaseTranslation): else: # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( - messages=messages, + data=data, responses=guardrailed_texts, scanned=scanned, ) + self._apply_guardrail_tool_calls_to_input( + messages=messages, + scanned_tool_calls=scanned_tool_calls, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + ) verbose_proxy_logger.debug("Anthropic Messages: Processed input messages: %s", messages) @@ -598,6 +710,19 @@ class AnthropicMessagesHandler(BaseTranslation): hoisted: Final = probe.get("messages") or [] # mutable-ok: API message payload return hoisted[0] if hoisted else None + @staticmethod + def _extract_top_level_system_text(hoisted_system_message: AllMessageValues) -> tuple[ScannedText, ...]: + content: Final = hoisted_system_message.get("content") + if isinstance(content, str): + return (ScannedText(content, SystemStringTarget()),) + if not isinstance(content, list): + return () + return tuple( + ScannedText(text_str, SystemBlockTextTarget(block_idx)) + for block_idx, block in enumerate(content) + if isinstance(block, dict) and isinstance(text_str := block.get("text"), str) and text_str + ) + @staticmethod def _openai_system_message_to_anthropic( message: Mapping[str, object], @@ -852,9 +977,25 @@ class AnthropicMessagesHandler(BaseTranslation): for content_idx, content_item in enumerate(content) if isinstance(content_item, dict) ) + tool_use_blocks: Final = ( + () + if scan_only_tool_results + else tuple( + (content_idx, content_item) + for content_idx, content_item in enumerate(content) + if isinstance(content_item, dict) and _is_client_tool_use(content_item) + ) + ) return ExtractedInput( scanned=tuple(item for block in blocks for item in block.scanned), images=tuple(image for block in blocks for image in block.images), + tool_calls=tuple( + ScannedToolCall( + tool_call=AnthropicConfig.convert_tool_use_to_openai_format(content_item, tool_call_idx), + target=ToolUseInputTarget(msg_idx, content_idx), + ) + for tool_call_idx, (content_idx, content_item) in enumerate(tool_use_blocks) + ), ) @classmethod @@ -940,43 +1081,48 @@ class AnthropicMessagesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - messages: Sequence[_WritableMessage], + data: dict, responses: list[str], scanned: tuple[ScannedText, ...], ) -> None: """ - Apply guardrail responses back to input messages. + Apply guardrail responses back to the top-level system prompt and the input messages. """ + messages: Final[Sequence[_WritableMessage]] = data.get("messages") or () for item, guardrail_response in zip(scanned, responses): - target = item.target - message = messages[target.msg_idx] - content = message.get("content", None) - if content is None: - continue - - match target: - case MessageContentTarget(): - if isinstance(content, str): - message["content"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ContentBlockTextTarget(content_idx=content_idx): - if isinstance(content, list): - content[content_idx]["text"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ToolResultStringTarget(content_idx=content_idx): - if isinstance(content, list): - content[content_idx]["content"] = ( - guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place - ) - case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx): - if isinstance(content, list): - content[content_idx]["content"][block_idx]["text"] = ( + match item.target: + case SystemStringTarget(): + if isinstance(data.get("system"), str): + data["system"] = ( guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place ) + case SystemBlockTextTarget(block_idx=block_idx): + _write_back_system_block(data.get("system"), block_idx, guardrail_response) + case ( + MessageContentTarget() + | ContentBlockTextTarget() + | ToolResultStringTarget() + | ToolResultBlockTextTarget() as message_target + ): + _write_back_message_text(messages[message_target.msg_idx], message_target, guardrail_response) case _: - assert_never(target) + assert_never(item.target) + + @staticmethod + def _apply_guardrail_tool_calls_to_input( + messages: Sequence[_WritableMessage], + scanned_tool_calls: tuple[ScannedToolCall, ...], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + returned_tool_calls: object, + ) -> None: + post_guardrail_tool_calls: Final = _tool_call_shapes( + returned_tool_calls + if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + else tuple(item.tool_call for item in scanned_tool_calls) + ) + for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls): + if before != after: + _write_back_tool_use(messages[item.target.msg_idx], item.target, after) async def process_output_response( self, diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 9fe56f4dc65..4b47b7d8c4a 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -635,14 +635,19 @@ class TestAnthropicMessagesHandlerInputProcessing: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) assert guardrail.inputs is not None - assert guardrail.inputs["texts"] == ["safe text", "prohibited correction"] + assert guardrail.inputs["texts"] == [ + "trusted top-level system prompt", + "safe text", + "prohibited correction", + ] structured = guardrail.inputs["structured_messages"] assert [m["role"] for m in structured] == ["system", "user", "system"] assert structured[0]["content"] == "trusted top-level system prompt" + assert data["system"] == "trusted top-level system prompt" assert data["messages"][1]["content"] == "[MASKED]" @pytest.mark.asyncio - async def test_bedrock_masking_slice_is_unavailable_when_top_level_system_is_included( + async def test_bedrock_masking_slice_lines_up_when_top_level_system_is_included( self, ): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( @@ -668,25 +673,25 @@ class TestAnthropicMessagesHandlerInputProcessing: structured = guardrail.inputs["structured_messages"] bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1") - assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) + 1 + assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) latest_user_index = bedrock._find_latest_message_index(structured, target_role="user") - assert ( - bedrock._locate_message_texts_slice( - structured_messages=structured, - target_index=latest_user_index, - texts=texts, - ) - is None - ) - assert ( - bedrock._merge_masked_texts( - masked_texts=["{MASKED}"], - texts=texts, - scanned_slice=None, - scanned_role_subset=True, - ) - == texts + scanned_slice = bedrock._locate_message_texts_slice( + structured_messages=structured, + target_index=latest_user_index, + texts=texts, ) + assert scanned_slice == (3, 1) + assert bedrock._merge_masked_texts( + masked_texts=["{MASKED}"], + texts=texts, + scanned_slice=scanned_slice, + scanned_role_subset=True, + ) == [ + "trusted top-level system prompt", + "safe text", + "prohibited correction", + "{MASKED}", + ] @pytest.mark.asyncio @pytest.mark.parametrize("skip_system_message_in_guardrail", [True, None]) @@ -1611,7 +1616,8 @@ class TestAnthropicMessagesIncrementalScan: ) assert mock_api.call_count == 1 assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [ - "What is the capital of France?" + "You are a helpful geography assistant.", + "What is the capital of France?", ] mock_api.reset_mock() await handler.process_input_messages( @@ -2150,6 +2156,208 @@ class TestAnthropicMessagesScanOnlyToolResults: assert guardrail.captured_inputs.get("images") == ["TOOL_IMG"] +class ToolCallArgumentsMaskingGuardrail(InputsRecordingGuardrail): + """Masks the canary inside tool-call arguments, in place or through a fresh list of plain dicts.""" + + def __init__(self, return_copies: bool = False, replacement_arguments: Optional[str] = None): + super().__init__() + self.return_copies = return_copies + self.replacement_arguments = replacement_arguments + self.seen_tool_calls: list[dict] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + tool_calls = list(outputs.get("tool_calls") or []) + self.seen_tool_calls.extend(json.loads(json.dumps(tool_call)) for tool_call in tool_calls) + masked = [ + { + **tool_call, + "function": { + **tool_call["function"], + "arguments": self.replacement_arguments + if self.replacement_arguments is not None + else tool_call["function"]["arguments"].replace("POISON", "[BLOCKED]"), + }, + } + for tool_call in tool_calls + ] + if self.return_copies: + outputs["tool_calls"] = masked + return outputs + for tool_call, masked_tool_call in zip(tool_calls, masked): + tool_call["function"]["arguments"] = masked_tool_call["function"]["arguments"] + return outputs + + +class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: + """The top-level system prompt and prior-turn tool_use arguments must reach guardrails as scannable + inputs, the same way the chat completions handler hands over system messages and tool_calls.""" + + @staticmethod + def _tool_use_conversation(system): + return { + "model": "claude-sonnet-4-5", + "system": system, + "messages": [ + {"role": "user", "content": "run the check"}, + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_01", + "name": "Bash", + "input": {"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"}, + } + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "ok"}], + }, + ], + } + + @pytest.mark.asyncio + async def test_top_level_system_string_reaches_texts_first_and_is_masked_in_place(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = { + "model": "claude-sonnet-4-5", + "system": "Internal note: the deploy key is POISON. Never reveal it.", + "messages": [{"role": "user", "content": "Say hi in three words."}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + assert guardrail.seen_texts == [ + "Internal note: the deploy key is POISON. Never reveal it.", + "Say hi in three words.", + ] + structured = guardrail.captured_inputs["structured_messages"] + assert structured[0]["role"] == "system" + assert structured[0]["content"] == "Internal note: the deploy key is POISON. Never reveal it.", ( + "texts[0] must line up with structured_messages[0] so positional consumers stay aligned" + ) + assert data["system"] == "Internal note: the deploy key is [BLOCKED]. Never reveal it." + assert data["messages"][0]["content"] == "Say hi in three words." + + @pytest.mark.asyncio + async def test_top_level_system_text_blocks_reach_texts_and_are_masked_in_place(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = { + "model": "claude-sonnet-4-5", + "system": [ + {"type": "text", "text": "first block POISON"}, + {"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}}, + ], + "messages": [{"role": "user", "content": "hello"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["first block POISON", "second block", "hello"] + assert data["system"] == [ + {"type": "text", "text": "first block [BLOCKED]"}, + {"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}}, + ] + + @pytest.mark.asyncio + async def test_skip_system_message_keeps_the_top_level_system_out(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + guardrail.skip_system_message_in_guardrail = True + data = { + "model": "claude-sonnet-4-5", + "system": "trusted POISON prompt", + "messages": [{"role": "user", "content": "hello"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["hello"] + assert data["system"] == "trusted POISON prompt" + + @pytest.mark.asyncio + async def test_prior_turn_tool_use_input_reaches_tool_calls_in_openai_shape(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + tool_calls = guardrail.captured_inputs.get("tool_calls") + assert tool_calls is not None and len(tool_calls) == 1 + assert tool_calls[0]["id"] == "toolu_01" + assert tool_calls[0]["type"] == "function" + assert tool_calls[0]["function"]["name"] == "Bash" + assert json.loads(tool_calls[0]["function"]["arguments"]) == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + }, "a guardrail that leaves tool_calls alone must leave the tool_use input alone" + + @pytest.mark.asyncio + @pytest.mark.parametrize("return_copies", [False, True]) + async def test_masked_tool_call_arguments_write_back_into_the_tool_use_input(self, return_copies: bool): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(return_copies=return_copies) + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [tool_call["function"]["name"] for tool_call in guardrail.seen_tool_calls] == ["Bash"] + tool_use = data["messages"][1]["content"][0] + assert tool_use == { + "type": "tool_use", + "id": "toolu_01", + "name": "Bash", + "input": {"cmd": "AWS_ACCESS_KEY_ID=[BLOCKED] aws sts get-caller-identity"}, + } + assert data["messages"][2]["content"][0]["tool_use_id"] == "toolu_01" + + @pytest.mark.asyncio + async def test_non_json_rewritten_arguments_keep_the_tool_use_input(self): + handler = AnthropicMessagesHandler() + guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]") + data = self._tool_use_conversation(system="You are a careful agent harness.") + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + + @pytest.mark.asyncio + async def test_scan_only_tool_results_keeps_system_and_tool_use_out(self): + handler = AnthropicMessagesHandler() + guardrail = InputsRecordingGuardrail() + guardrail.scan_only_tool_results = True + data = self._tool_use_conversation(system="trusted POISON prompt") + data["messages"][2]["content"][0]["content"] = "fetched POISON page" + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["fetched POISON page"] + assert guardrail.captured_inputs is not None + assert guardrail.captured_inputs.get("tool_calls") is None + assert data["system"] == "trusted POISON prompt" + assert data["messages"][1]["content"][0]["input"] == { + "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" + } + assert data["messages"][2]["content"][0]["content"] == "fetched [BLOCKED] page" + + class TestStructuredWriteBackKeepsToolResults: """A guardrail rewrite must never leave a tool_use without its tool_result (Claude Code ToolSearch, LIT-6103).""" From c0c0c9a9ebc7443b1724c22e3b6b856aac1f750d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:05:10 -0700 Subject: [PATCH 034/107] fix(sdk): keep body and proxy headers on BadRequestError mapped from a litellm_proxy 400 The generic 400 branch of the OpenAI exception mapper dropped the wire body and no branch carried the response headers, so an application calling a LiteLLM proxy through a litellm_proxy/ model could not tell a guardrail block from any other failure without walking __cause__. BadRequestError now takes headers, filled for a litellm_proxy upstream, and the generic branch passes the body. The proxy edge treats the literal "None" type and param an older proxy sends as absent and stops forwarding an upstream proxy's date and server headers. --- litellm/constants.py | 6 +- litellm/exceptions.py | 4 +- .../exception_mapping_utils.py | 16 +++++ .../common_utils/openai_error_payload.py | 6 +- .../test_exception_mapping_utils.py | 68 +++++++++++++++++++ .../common_utils/test_openai_error_payload.py | 18 +++++ .../proxy/test_common_request_processing.py | 15 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 8 files changed, 129 insertions(+), 6 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5751e6e46af..caf1aff6792 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1964,7 +1964,11 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( } ) -UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS +ORIGIN_SERVER_HEADERS: Final[frozenset[str]] = frozenset({"date", "server"}) + +UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = ( + HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS | ORIGIN_SERVER_HEADERS +) # A retrieved response replays the usage of the call that created it, so pricing these # read/management routes like inference bills the same tokens twice. diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 3f22a4b2dcd..3d9e26e450b 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -10,7 +10,7 @@ ## LiteLLM versions of the OpenAI Exception Types import enum -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final import httpx @@ -226,6 +226,7 @@ class BadRequestError(openai.BadRequestError): max_retries: int | None = None, num_retries: int | None = None, body: dict | None = None, + headers: Mapping[str, str] | None = None, ): self.status_code = 400 self.message = f"litellm.BadRequestError: {message}" @@ -234,6 +235,7 @@ class BadRequestError(openai.BadRequestError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries + self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None # Use response if it's a valid httpx.Response with a request, otherwise use minimal error response # Note: We check _request (not .request property) to avoid RuntimeError when _request is None if ( diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 82708d412c9..fd1ba666887 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1,6 +1,7 @@ import json import re import traceback +from collections.abc import Mapping from typing import Any, Final, Protocol, cast import httpx @@ -254,6 +255,15 @@ class _ProviderHTTPException(Protocol): llm_provider: str +def _litellm_proxy_response_headers( + original_exception: _ProviderHTTPException, custom_llm_provider: str +) -> Mapping[str, str] | None: + if custom_llm_provider != "litellm_proxy": + return None + headers: Final = getattr(original_exception, "headers", None) + return headers if isinstance(headers, Mapping) else None + + def _map_openai_exception( *, model: str, @@ -264,6 +274,7 @@ def _map_openai_exception( exception_provider: str, extra_information: str, ) -> None: + upstream_headers: Final = _litellm_proxy_response_headers(original_exception, custom_llm_provider) # custom_llm_provider is openai, make it OpenAI message = get_error_message(error_obj=original_exception) if message is None: @@ -348,6 +359,7 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str: raise BadRequestError( @@ -357,6 +369,7 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif ( "Web server is returning an unknown error" in error_str @@ -404,6 +417,8 @@ def _map_openai_exception( model=model, response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif original_exception.status_code == 401: raise AuthenticationError( @@ -436,6 +451,7 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif original_exception.status_code == 429: raise RateLimitError( diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index 89f735ee8b6..90b3c998247 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -8,6 +8,8 @@ from typing import Final from fastapi import status +_STRINGIFIED_NONE: Final = "None" + _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( { status.HTTP_401_UNAUTHORIZED: "authentication_error", @@ -35,7 +37,7 @@ def openai_error_type(exc: object, status_code: int) -> str: """OpenAI types ``error.type`` as a required string, so an exception carrying none falls back to the type its status code stands for.""" carried: Final = attribute_of(exc, "type") - if isinstance(carried, str): + if isinstance(carried, str) and carried != _STRINGIFIED_NONE: return carried mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) if mapped is not None: @@ -49,4 +51,4 @@ def openai_error_param(exc: object) -> str | None: """OpenAI types ``error.param`` as nullable, so an exception carrying none serializes as JSON ``null``.""" carried: Final = attribute_of(exc, "param") - return carried if isinstance(carried, str) else None + return carried if isinstance(carried, str) and carried != _STRINGIFIED_NONE else None diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 42d3df76902..ea6ac17ad45 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1409,3 +1409,71 @@ def test_bedrock_timeout_mapping_keeps_retry_after_readable(status_code): exception_headers = _get_response_headers(original_exception=exc_info.value) assert exception_headers is not None assert litellm.utils._get_retry_after_from_exception_header(response_headers=exception_headers) == 7 + + +_GUARDRAIL_BLOCK_ERROR = { + "message": "Content blocked: secret_project_codename pattern detected", + "param": "None", + "code": "400", + "provider_specific_fields": { + "error": "Content blocked: secret_project_codename pattern detected", + "pattern": "secret_project_codename", + "guardrail_name": "block-secret-project", + "guardrail_mode": "pre_call", + }, +} + + +def _openai_handler_error(error_type: str, headers: dict[str, str]) -> OpenAIError: + """What litellm/llms/openai/openai.py raises after the openai SDK rejects a 400: + the SDK's str() carries the wire body, and the handler copies headers and body over.""" + wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type} + wire = httpx.Response( + status_code=400, + headers=headers, + json={"error": wire_error}, + request=httpx.Request("POST", "http://localhost:4000/v1/chat/completions"), + ) + return OpenAIError( + status_code=400, + message=f"Error code: 400 - {{'error': {wire_error}}}", + headers=wire.headers, + body=wire_error, + ) + + +@pytest.mark.parametrize("error_type", ["None", "invalid_request_error"]) +def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str): + """An SDK caller behind a proxy tells a guardrail block from any other 400 by the body's + provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError + must carry both whichever error.type the proxy version on the other end emits.""" + proxy_headers = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} + + with pytest.raises(litellm.BadRequestError) as exc_info: + exception_type( + model="claude-haiku-4-5", + original_exception=_openai_handler_error(error_type, proxy_headers), + custom_llm_provider="litellm_proxy", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" + assert exc_info.value.body["type"] == error_type + assert proxy_headers.items() <= exc_info.value.headers.items() + + +def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): + """A vendor's own response headers stay on e.response the way every other mapped provider + error keeps them; only a LiteLLM proxy upstream puts headers on e.headers.""" + with pytest.raises(litellm.BadRequestError) as exc_info: + exception_type( + model="gpt-5.4-mini", + original_exception=_openai_handler_error("vendor_specific_error", {"openai-organization": "org-1"}), + custom_llm_provider="openai", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.body["type"] == "vendor_specific_error" + assert exc_info.value.headers is None diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index 8b653ddfb71..db9fe3a2253 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -143,3 +143,21 @@ def test_a_status_carried_by_an_exception_drives_the_type_it_reports(): exc = HTTPException(status_code=403, detail="blocked by policy") assert openai_error_type(exc, error_status_code(exc, 400)) == "permission_error" + + +def test_the_stringified_none_an_older_upstream_proxy_sent_is_treated_as_absent(): + """A proxy fronting a proxy older than 1.102 receives {"type": "None", "param": "None"} on + the wire; the SDK now keeps that body on the mapped exception, and re-emitting the literal + is the exact bug this module exists to stop.""" + from litellm.exceptions import BadRequestError + + carried = BadRequestError( + message="Content blocked", + model="claude-haiku-4-5", + llm_provider="litellm_proxy", + body={"message": "Content blocked", "type": "None", "param": "None", "code": "400"}, + ) + + assert carried.type == "None" + assert openai_error_type(carried, 400) == "invalid_request_error" + assert openai_error_param(carried) is None diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 69e89d1c604..3cb58c44354 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3985,6 +3985,21 @@ class TestHandleLLMApiExceptionFramingHeaders: assert proxy_exc.headers["x-custom-safe"] == "1" assert proxy_exc.headers["x-request-id"] == "abc-123" + async def test_strips_the_date_and_server_headers_of_an_upstream_litellm_proxy(self): + """A proxy fronting another LiteLLM proxy gets the upstream's date and server + on the mapped exception; forwarding them would duplicate the Date header + uvicorn adds to every response and leak the upstream server identity.""" + exc = litellm.BadRequestError( + message="Content blocked", + llm_provider="litellm_proxy", + model="claude-haiku-4-5", + headers={"date": "Sun, 13 Sep 2026 08:43:51 GMT", "server": "uvicorn", "x-request-id": "abc-123"}, + ) + proxy_exc = await self._invoke(exc) + assert "date" not in proxy_exc.headers + assert "server" not in proxy_exc.headers + assert proxy_exc.headers["x-request-id"] == "abc-123" + class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From e732a484f6b427d6513bc39885fdc8485ba29e7c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:11:59 -0700 Subject: [PATCH 035/107] fix(responses): hoist Codex additional_tools input items into the chat bridge tools --- .../responses/transformation.py | 66 +++--------- litellm/responses/additional_tools.py | 65 +++++++++++ .../custom_tools.py | 28 +++-- .../handler.py | 22 ++-- .../transformation.py | 22 ++-- .../test_handler.py | 102 ++++++++++++++++++ .../test_litellm_completion_responses.py | 91 ++++++++++++++++ .../responses/test_additional_tools.py | 48 +++++++++ .../responses/test_custom_tool_call.py | 18 ++++ 9 files changed, 387 insertions(+), 75 deletions(-) create mode 100644 litellm/responses/additional_tools.py create mode 100644 tests/test_litellm/responses/test_additional_tools.py diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 53a3e634adf..95375399033 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -17,7 +17,7 @@ BaseAWSLLM._sign_request after the request body is finalized. import json from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict import httpx from typing_extensions import ReadOnly, TypedDict @@ -32,6 +32,7 @@ from litellm.llms.bedrock_mantle.common_utils import ( BedrockMantleAuthMixin, ) from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( ResponseInputParam, @@ -58,8 +59,6 @@ _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset( _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"}) _BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES: Final = frozenset({"auto"}) -_CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools" - _CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message" _CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction" _CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call" @@ -233,62 +232,29 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI litellm_params: GenericLiteLLMParams, headers: dict, ) -> dict: - remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input) - normalized_input: Final = self._normalize_codex_input_items(remaining_input) - request_params: Final = ( - { - **response_api_optional_request_params, - "tools": [ - *(response_api_optional_request_params.get("tools") or []), - *hoisted_tools, - ], - } - if hoisted_tools - else response_api_optional_request_params + params: Final = cast( # cast-ok: the base signature leaves the params dict untyped + "ResponsesAPIOptionalRequestParams", response_api_optional_request_params ) + hoisted: Final = hoist_additional_tools(input, params.get("tools")) + normalized_input: Final = self._normalize_codex_input_items(hoisted.input) return super().transform_responses_api_request( model=model, input=normalized_input, - response_api_optional_request_params=request_params, + response_api_optional_request_params=self._params_with_hoisted_tools(params, hoisted), litellm_params=litellm_params, headers=headers, ) - @staticmethod - def _is_codex_additional_tools_item(item: Any) -> bool: - return isinstance(item, dict) and item.get("type") == _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE - - @staticmethod - def _tools_of_additional_tools_item(item: "dict[str, Any]") -> "list[Any]": - tools: Final = item.get("tools") - return tools if isinstance(tools, list) else [] - @classmethod - def _hoist_codex_additional_tools( - cls, - input: "str | ResponseInputParam", - ) -> "tuple[str | ResponseInputParam, list[Any]]": - """Codex's "responses lite" wire mode ships tool definitions inside - `input` as {"type": "additional_tools", "role": "developer", - "tools": [...]} items. api.openai.com accepts that item type; Mantle - rejects the whole request with 400 "Invalid 'input': value did not - match any expected variant" but accepts the same tools at the top - level, so move them there and strip the items from `input`. - """ - if not isinstance(input, list): - return input, [] - additional_tools_items: Final = [item for item in input if cls._is_codex_additional_tools_item(item)] - if not additional_tools_items: - return input, [] - remaining_input: Final = [item for item in input if not cls._is_codex_additional_tools_item(item)] - hoisted_tools = [tool for item in additional_tools_items for tool in cls._tools_of_additional_tools_item(item)] - verbose_logger.debug( - "Bedrock Mantle Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) " - "into the top-level tools param (Mantle rejects that input item type).", - len(hoisted_tools), - len(additional_tools_items), - ) - return remaining_input, cls._filter_unsupported_tools(hoisted_tools) + def _params_with_hoisted_tools( + cls, params: Mapping[str, object], hoisted: HoistedAdditionalTools + ) -> dict[str, object]: + if not hoisted.hoisted: + return dict(params) + supported_tools: Final = cls._filter_unsupported_tools(list(hoisted.tools)) + if supported_tools: + return {**params, "tools": supported_tools} + return {key: value for key, value in params.items() if key != "tools"} @staticmethod def _agent_message_text(item: "Mapping[str, object]") -> str: diff --git a/litellm/responses/additional_tools.py b/litellm/responses/additional_tools.py new file mode 100644 index 00000000000..ea0d7af350c --- /dev/null +++ b/litellm/responses/additional_tools.py @@ -0,0 +1,65 @@ +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Final, cast # noqa: TID251 # validating the openai tool union strips vendor keys from raw tools + +from pydantic import BaseModel, ValidationError + +from litellm._logging import verbose_logger +from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam + +ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools" + + +class _InputItemType(BaseModel): + type: str = "" + + +class _AdditionalToolsItem(BaseModel): + tools: tuple[dict[str, object], ...] = () + + +@dataclass(frozen=True, slots=True) +class HoistedAdditionalTools: + input: str | ResponseInputParam + tools: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...] + hoisted: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...] + + +def _is_additional_tools_item(item: object) -> bool: + try: + return _InputItemType.model_validate(item).type == ADDITIONAL_TOOLS_INPUT_ITEM_TYPE + except ValidationError: + return False + + +def _tools_of_item(item: object) -> tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]: + try: + parsed: Final = _AdditionalToolsItem.model_validate(item) + except ValidationError: + return () + return tuple( + cast( + "ALL_RESPONSES_API_TOOL_PARAMS", tool + ) # cast-ok: nested tools carry the same raw tool JSON as top-level tools + for tool in parsed.tools + ) + + +def hoist_additional_tools( + input: str | ResponseInputParam, + tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None, +) -> HoistedAdditionalTools: + existing: Final = tuple(tools or ()) + if isinstance(input, str): + return HoistedAdditionalTools(input=input, tools=existing, hoisted=()) + items: Final = tuple(item for item in input if _is_additional_tools_item(item)) + if not items: + return HoistedAdditionalTools(input=input, tools=existing, hoisted=()) + hoisted: Final = tuple(tool for item in items for tool in _tools_of_item(item)) + verbose_logger.debug( + "Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) into the top-level tools param.", + len(hoisted), + len(items), + ) + remaining_input: Final = [item for item in input if not _is_additional_tools_item(item)] + return HoistedAdditionalTools(input=remaining_input, tools=(*existing, *hoisted), hoisted=hoisted) diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index 4aa489d9e50..038964055c3 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -39,15 +39,27 @@ def openai_shaped_tool_call_item_id(item_type: str, tool_id: str) -> str: return f"{prefix}_{tool_id}" +class _ToolNameFields(BaseModel): + type: str = "" + name: str = "" + tools: tuple[object, ...] = () + + +def _custom_tool_names_of(tool: object) -> tuple[str, ...]: + try: + parsed: Final = _ToolNameFields.model_validate(tool) + except ValidationError: + return () + if parsed.type == "custom": + return (parsed.name,) if parsed.name else () + if parsed.type != "namespace": + return () + return tuple(name for nested_tool in parsed.tools for name in _custom_tool_names_of(nested_tool)) + + def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]: - """Extract names of tools originally defined as ``type: "custom"``.""" - if not tools: - return set() - names: Final[set[str]] = set() - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "custom" and "name" in tool: - names.add(tool["name"]) - return names + """Extract names of tools defined as ``type: "custom"``, at the top level or inside a ``namespace`` tool.""" + return {name for tool in tools or () for name in _custom_tool_names_of(tool)} def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool: diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index a0e8cd278e6..505b5b09433 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -6,6 +6,7 @@ from collections.abc import Coroutine, Mapping from typing import Final import litellm +from litellm.responses.additional_tools import hoist_additional_tools from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) @@ -37,11 +38,16 @@ class LiteLLMCompletionTransformationHandler: | BaseResponsesAPIStreamingIterator | Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator] ): + hoisted: Final = hoist_additional_tools(input, responses_api_request.get("tools")) + bridged_input: Final = hoisted.input + bridged_request: Final[ResponsesAPIOptionalRequestParams] = ( + {**responses_api_request, "tools": list(hoisted.tools)} if hoisted.hoisted else responses_api_request + ) litellm_completion_request: Final[dict] = ( LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model=model, - input=input, - responses_api_request=responses_api_request, + input=bridged_input, + responses_api_request=bridged_request, custom_llm_provider=custom_llm_provider, stream=stream, extra_headers=extra_headers, @@ -52,8 +58,8 @@ class LiteLLMCompletionTransformationHandler: if _is_async: return self.async_response_api_handler( litellm_completion_request=litellm_completion_request, - request_input=input, - responses_api_request=responses_api_request, + request_input=bridged_input, + responses_api_request=bridged_request, **kwargs, ) @@ -70,8 +76,8 @@ class LiteLLMCompletionTransformationHandler: responses_api_response: Final[ResponsesAPIResponse] = ( LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( chat_completion_response=litellm_completion_response, - request_input=input, - responses_api_request=responses_api_request, + request_input=bridged_input, + responses_api_request=bridged_request, ) ) @@ -81,8 +87,8 @@ class LiteLLMCompletionTransformationHandler: return LiteLLMCompletionStreamingIterator( model=model, litellm_custom_stream_wrapper=litellm_completion_response, - request_input=input, - responses_api_request=responses_api_request, + request_input=bridged_input, + responses_api_request=bridged_request, custom_llm_provider=custom_llm_provider, litellm_metadata=kwargs.get("litellm_metadata", {}), ) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index e8aacac9e67..27756b405ec 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1890,9 +1890,21 @@ class LiteLLMCompletionResponsesConfig: namespace_tool: NamespaceTool, nested: bool, ) -> ChatCompletionToolParam | None: - if nested and namespace_tool.get("type") != "function": + tool_type: Final = namespace_tool.get("type") + if nested and tool_type not in ("function", "custom"): return None + raw_description: Final = str(namespace_tool.get("description") or "") + description: Final = ( + f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}" + if nested and namespace_description and raw_description + else namespace_description + if nested and namespace_description + else raw_description + ) + if nested and tool_type == "custom": + return convert_custom_tool_to_function_tool({**namespace_tool, "description": description}) + raw_parameters: Final = namespace_tool.get("parameters") parameters: Final = ( MappingProxyType(raw_parameters) if isinstance(raw_parameters, Mapping) else MappingProxyType({}) @@ -1901,14 +1913,6 @@ class LiteLLMCompletionResponsesConfig: parameters if parameters and "type" in parameters else MappingProxyType({**parameters, "type": "object"}) ) tool_name: Final = str(namespace_tool.get("name") or "") - raw_description: Final = str(namespace_tool.get("description") or "") - description: Final = ( - f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}" - if nested and namespace_description and raw_description - else namespace_description - if nested and namespace_description - else raw_description - ) chat_tool_name: Final = f"{namespace}__{tool_name}" if nested else tool_name function: Final = ChatCompletionToolParamFunctionChunk( name=chat_tool_name, diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py index 2cfec6a1844..b78dabbfe48 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py @@ -68,3 +68,105 @@ async def test_async_fallback_tags_skip_responses_api_bridge(): await coro assert captured.get("_skip_responses_api_bridge") is True + + +_CODEX_ADDITIONAL_TOOLS_ITEM = { + "type": "additional_tools", + "id": "at_codex", + "role": "developer", + "tools": [ + { + "type": "namespace", + "name": "functions", + "description": "", + "tools": [ + { + "type": "custom", + "name": "exec", + "description": "Runs a shell command.", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}, + }, + { + "type": "function", + "name": "wait", + "description": "Waits for a background command.", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}}, + }, + ], + } + ], +} +_CODEX_INPUT = [_CODEX_ADDITIONAL_TOOLS_ITEM, {"type": "message", "role": "user", "content": "Run ls"}] + + +def test_sync_fallback_hoists_additional_tools_input_items_into_chat_tools(): + handler = LiteLLMCompletionTransformationHandler() + captured: dict = {} + + def fake_completion(**kwargs): + captured.update(kwargs) + raise _StopForwarding() + + with patch("litellm.completion", fake_completion): # test-quality-ok: no DI seam; the file stubs this same boundary + with pytest.raises(_StopForwarding): + handler.response_api_handler( + model="bedrock/us.openai.gpt-5.6", + input=_CODEX_INPUT, + responses_api_request={}, + custom_llm_provider="bedrock", + _is_async=False, + ) + + assert [message["role"] for message in captured["messages"]] == ["user"] + functions_by_name = {tool["function"]["name"]: tool["function"] for tool in captured["tools"]} + assert set(functions_by_name) == {"exec", "functions__wait"} + assert set(functions_by_name["exec"]["parameters"]["properties"]) == {"content"} + + +@pytest.mark.asyncio +async def test_async_fallback_returns_hoisted_nested_custom_tool_call_as_custom_tool_call(): + from litellm.responses.litellm_completion_transformation.transformation import TOOL_CALLS_CACHE + from litellm.types.utils import ChatCompletionMessageToolCall, Choices, Function, Message, ModelResponse + + handler = LiteLLMCompletionTransformationHandler() + tool_call_id = "call_exec_hoisted" + + async def fake_acompletion(**kwargs): + return ModelResponse( + id="chatcmpl-exec", + created=1, + model="us.openai.gpt-5.6", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id=tool_call_id, + type="function", + function=Function(name="exec", arguments='{"content": "ls"}'), + ) + ], + ), + ) + ], + ) + + try: + with patch("litellm.acompletion", fake_acompletion): # test-quality-ok: no DI seam; file stubs this boundary + response = await handler.response_api_handler( + model="bedrock/us.openai.gpt-5.6", + input=_CODEX_INPUT, + responses_api_request={}, + custom_llm_provider="bedrock", + _is_async=True, + ) + finally: + TOOL_CALLS_CACHE.delete_cache(key=tool_call_id) + + tool_calls = [(item.type, item.name, item.input) for item in response.output if item.type == "custom_tool_call"] + assert tool_calls == [("custom_tool_call", "exec", "ls")] diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 3c78bbf79d7..1f497597a11 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2506,6 +2506,7 @@ class TestToolTransformation: "tools": [ "ignored", {"type": "namespace", "name": "ignored"}, + {"type": "web_search", "name": "ignored"}, { "type": "function", "name": "spawn_agent", @@ -2527,6 +2528,36 @@ class TestToolTransformation: "type": "object", } + def test_transform_nested_namespace_custom_tool_becomes_a_content_function_under_its_short_name(self): + namespace_tool = { + "type": "namespace", + "name": "functions", + "description": "Codex shell tools.", + "tools": [ + { + "type": "custom", + "name": "exec", + "description": "Runs a shell command.", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}, + }, + ], + } + + result_tools, _ = ( + LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[namespace_tool] + ) + ) + + assert len(result_tools) == 1 + function = result_tools[0]["function"] + assert function["name"] == "exec" + assert function["description"].startswith("Codex shell tools.") + assert "Runs a shell command." in function["description"] + assert "start: /.+/" in function["description"] + assert function["parameters"]["required"] == ["content"] + assert function["parameters"]["properties"]["content"]["type"] == "string" + @pytest.mark.parametrize( "model, custom_llm_provider", [ @@ -3786,6 +3817,66 @@ class TestEnsureOutputItemContentPartAdded: assert added.item.name == "spawn_agent" assert added.item.namespace == "collaboration" + def test_streaming_nested_custom_tool_call_comes_back_as_custom_tool_call(self): + from litellm.responses.litellm_completion_transformation.custom_tools import extract_custom_tool_names + + iterator = self._make_iterator() + iterator.responses_api_request = { + "tools": [ + { + "type": "namespace", + "name": "functions", + "tools": [ + { + "type": "custom", + "name": "exec", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}, + } + ], + } + ] + } + iterator._custom_tool_names = extract_custom_tool_names(iterator.responses_api_request.get("tools")) + iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( + iterator.responses_api_request.get("tools") + ) + + iterator._queue_tool_call_delta_events( + [{"index": 0, "id": "call_exec", "function": {"name": "exec", "arguments": '{"content":"ls"}'}}] + ) + iterator._queue_final_tool_call_done_events( + ModelResponse( + id="chatcmpl-exec", + created=1, + model="us.openai.gpt-5.6", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_exec", + type="function", + function=Function(name="exec", arguments='{"content":"ls"}'), + ) + ], + ), + ) + ], + ) + ) + + added = iterator._pending_tool_events[0] + assert added.item.type == "custom_tool_call" + assert added.item.name == "exec" + done = iterator._pending_tool_events[-1] + assert done.item.type == "custom_tool_call" + assert done.item.input == "ls" + def test_streaming_unqualified_namespace_tool_calls_restore_namespace(self): """A unique nested tool name without the namespace still maps back.""" iterator = self._make_iterator() diff --git a/tests/test_litellm/responses/test_additional_tools.py b/tests/test_litellm/responses/test_additional_tools.py new file mode 100644 index 00000000000..bef3b27eacd --- /dev/null +++ b/tests/test_litellm/responses/test_additional_tools.py @@ -0,0 +1,48 @@ +from litellm.responses.additional_tools import hoist_additional_tools + +_EXEC_TOOL = {"type": "custom", "name": "exec", "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}} +_WAIT_TOOL = {"type": "function", "name": "wait", "parameters": {"type": "object", "properties": {}}} +_TOP_LEVEL_TOOL = {"type": "function", "name": "top_level", "parameters": {"type": "object", "properties": {}}} +_USER_MESSAGE = {"type": "message", "role": "user", "content": "Run ls"} + + +def test_string_input_passes_through_with_existing_tools(): + hoisted = hoist_additional_tools("hello", [_TOP_LEVEL_TOOL]) + + assert hoisted.input == "hello" + assert hoisted.tools == (_TOP_LEVEL_TOOL,) + assert hoisted.hoisted == () + + +def test_input_without_additional_tools_items_is_returned_untouched(): + request_input = [_USER_MESSAGE] + + hoisted = hoist_additional_tools(request_input, None) + + assert hoisted.input is request_input + assert hoisted.tools == () + assert hoisted.hoisted == () + + +def test_additional_tools_items_are_stripped_and_appended_after_top_level_tools_in_item_order(): + request_input = [ + {"type": "additional_tools", "id": "at_1", "role": "developer", "tools": [_EXEC_TOOL]}, + _USER_MESSAGE, + {"type": "additional_tools", "id": "at_2", "role": "developer", "tools": [_WAIT_TOOL]}, + ] + + hoisted = hoist_additional_tools(request_input, [_TOP_LEVEL_TOOL]) + + assert hoisted.input == [_USER_MESSAGE] + assert hoisted.tools == (_TOP_LEVEL_TOOL, _EXEC_TOOL, _WAIT_TOOL) + assert hoisted.hoisted == (_EXEC_TOOL, _WAIT_TOOL) + + +def test_additional_tools_item_without_a_tools_list_is_stripped_and_contributes_nothing(): + request_input = [{"type": "additional_tools", "id": "at_1", "role": "developer", "tools": "exec"}, _USER_MESSAGE] + + hoisted = hoist_additional_tools(request_input, None) + + assert hoisted.input == [_USER_MESSAGE] + assert hoisted.tools == () + assert hoisted.hoisted == () diff --git a/tests/test_litellm/responses/test_custom_tool_call.py b/tests/test_litellm/responses/test_custom_tool_call.py index e80301c3b2f..2ed71ee3ecf 100644 --- a/tests/test_litellm/responses/test_custom_tool_call.py +++ b/tests/test_litellm/responses/test_custom_tool_call.py @@ -55,6 +55,24 @@ class TestCustomToolUtilities: names = extract_custom_tool_names(tools) assert names == set() + def test_extract_custom_tool_names_walks_namespace_tools(self): + tools = [ + {"type": "function", "name": "regular_tool"}, + { + "type": "namespace", + "name": "functions", + "tools": [ + {"type": "custom", "name": "exec"}, + {"type": "function", "name": "wait"}, + "ignored", + ], + }, + {"type": "namespace", "name": "empty", "tools": "not-a-list"}, + ] + + names = extract_custom_tool_names(tools) + assert names == {"exec"} + def test_extract_custom_tool_names_none(self): """Test extraction with None input.""" names = extract_custom_tool_names(None) From f7e9277032651786d493c72225e571d28f40c136 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:13:17 -0700 Subject: [PATCH 036/107] fix(guardrails): validate rewritten tool_use arguments with a typed adapter --- .../chat/guardrail_translation/handler.py | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index eb16278c9bd..b438d168b52 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -20,6 +20,7 @@ from itertools import chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable +from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict, assert_never from litellm._logging import verbose_proxy_logger @@ -188,7 +189,7 @@ def _is_client_tool_use(block: Mapping[str, object]) -> bool: block.get("type") == "tool_use" and isinstance(block.get("id"), str) and isinstance(block.get("name"), str) - and isinstance(block.get("input"), Mapping) + and isinstance(block.get("input"), dict) ) @@ -229,22 +230,19 @@ def _write_back_message_text(message: _WritableMessage, target: MessageTextTarge assert_never(target) +_TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object]) + + def _write_back_tool_use(message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape) -> None: content: Final = message.get("content", None) block: Final = content[target.content_idx] if isinstance(content, list) else None if not isinstance(block, dict): return try: - rewritten_input: Final = json.loads(shape.arguments) - except json.JSONDecodeError: + rewritten_input: Final = _TOOL_USE_INPUT_ADAPTER.validate_json(shape.arguments) + except ValidationError: verbose_proxy_logger.warning( - "Anthropic Messages: guardrail returned non-JSON arguments for tool_use %s; keeping its input", - block.get("id"), - ) - return - if not isinstance(rewritten_input, dict): - verbose_proxy_logger.warning( - "Anthropic Messages: guardrail returned non-object arguments for tool_use %s; keeping its input", + "Anthropic Messages: guardrail returned arguments that are not a JSON object for tool_use %s; keeping its input", block.get("id"), ) return @@ -1113,11 +1111,11 @@ class AnthropicMessagesHandler(BaseTranslation): messages: Sequence[_WritableMessage], scanned_tool_calls: tuple[ScannedToolCall, ...], pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], - returned_tool_calls: object, + returned_tool_calls: Sequence[object] | None, ) -> None: post_guardrail_tool_calls: Final = _tool_call_shapes( returned_tool_calls - if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls) else tuple(item.tool_call for item in scanned_tool_calls) ) for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls): From 825e4f17e949439d37f922bb02c0356f2bdd0dc5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:23:34 -0700 Subject: [PATCH 037/107] fix(sdk): carry body and proxy headers on relayed litellm errors and content policy blocks too --- litellm/exceptions.py | 2 + .../exception_mapping_utils.py | 49 +++++++------- .../test_exception_mapping_utils.py | 65 +++++++++++++------ .../common_utils/test_openai_error_payload.py | 2 +- 4 files changed, 72 insertions(+), 46 deletions(-) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 3d9e26e450b..fdc2cc1f169 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -626,6 +626,7 @@ class ContentPolicyViolationError(BadRequestError): litellm_debug_info: str | None = None, provider_specific_fields: dict | None = None, body: dict | None = None, + headers: Mapping[str, str] | None = None, ): self.status_code = 400 self.message = f"litellm.ContentPolicyViolationError: {message}" @@ -640,6 +641,7 @@ class ContentPolicyViolationError(BadRequestError): response=response, litellm_debug_info=self.litellm_debug_info, body=body, + headers=headers, ) # Call the base class constructor with the parameters it needs def __str__(self): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index fd1ba666887..36b53a26c99 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1,3 +1,4 @@ +import inspect import json import re import traceback @@ -203,11 +204,18 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None return _response_headers +def _accepted_init_kwargs(exception_class: type[Exception], candidates: Mapping[str, object]) -> Mapping[str, object]: + accepted: Final = inspect.signature(exception_class).parameters + return {name: value for name, value in candidates.items() if name in accepted} + + def extract_and_raise_litellm_exception( response: Any | None, error_str: str, model: str, custom_llm_provider: str, + body: object | None = None, + headers: Mapping[str, str] | None = None, ): """ Covers scenario where litellm sdk calling proxy. @@ -217,32 +225,19 @@ def extract_and_raise_litellm_exception( Relevant Issue: https://github.com/BerriAI/litellm/issues/7259 """ pattern: Final = r"litellm\.\w+Error" - - # Search for the exception in the error string match: Final = re.search(pattern, error_str) - - # Extract the exception if found - if match: - exception_name = match.group(0) - exception_name = exception_name.strip().replace("litellm.", "") - raised_exception_obj: Final = getattr(litellm, exception_name, None) - if raised_exception_obj: - # Try with response parameter first, fall back to without it - # Some exceptions (e.g., APIConnectionError) don't accept response param - try: - raise raised_exception_obj( - message=error_str, - llm_provider=custom_llm_provider, - model=model, - response=response, - ) - except TypeError: - # Exception doesn't accept response parameter - raise raised_exception_obj( - message=error_str, - llm_provider=custom_llm_provider, - model=model, - ) + if match is None: + return + exception_name: Final = match.group(0).removeprefix("litellm.") + raised_exception_obj: Final = getattr(litellm, exception_name, None) + if not raised_exception_obj: + return + raise raised_exception_obj( + message=error_str, + llm_provider=custom_llm_provider, + model=model, + **_accepted_init_kwargs(raised_exception_obj, {"response": response, "body": body, "headers": headers}), + ) class _ProviderHTTPException(Protocol): @@ -339,6 +334,8 @@ def _map_openai_exception( model=model, response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str: helpful_message: Final = ( @@ -2443,6 +2440,8 @@ def exception_type( error_str=error_str, model=model, custom_llm_provider=custom_llm_provider, + body=getattr(original_exception, "body", None), + headers=_litellm_proxy_response_headers(mappable_exception, custom_llm_provider), ) if ( custom_llm_provider == "openai" diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index ea6ac17ad45..1ff2bbb9bdd 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1424,35 +1424,37 @@ _GUARDRAIL_BLOCK_ERROR = { } -def _openai_handler_error(error_type: str, headers: dict[str, str]) -> OpenAIError: - """What litellm/llms/openai/openai.py raises after the openai SDK rejects a 400: +def _openai_handler_error( + error_type: str, + headers: dict[str, str], + status_code: int = 400, + message: str = _GUARDRAIL_BLOCK_ERROR["message"], +) -> OpenAIError: + """What litellm/llms/openai/openai.py raises after the openai SDK rejects a request: the SDK's str() carries the wire body, and the handler copies headers and body over.""" - wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type} - wire = httpx.Response( - status_code=400, - headers=headers, - json={"error": wire_error}, - request=httpx.Request("POST", "http://localhost:4000/v1/chat/completions"), - ) + wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type, "code": str(status_code), "message": message} return OpenAIError( - status_code=400, - message=f"Error code: 400 - {{'error': {wire_error}}}", - headers=wire.headers, + status_code=status_code, + message=f"Error code: {status_code} - {{'error': {wire_error}}}", + headers=httpx.Headers(headers), body=wire_error, ) -@pytest.mark.parametrize("error_type", ["None", "invalid_request_error"]) -def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str): - """An SDK caller behind a proxy tells a guardrail block from any other 400 by the body's - provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError - must carry both whichever error.type the proxy version on the other end emits.""" - proxy_headers = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} +_PROXY_HEADERS = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} + +@pytest.mark.parametrize( + ("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)] +) +def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, status_code: int): + """An SDK caller behind a proxy tells a guardrail block from any other 4xx by the body's + provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError + must carry both whichever error.type and status the proxy version on the other end emits.""" with pytest.raises(litellm.BadRequestError) as exc_info: exception_type( model="claude-haiku-4-5", - original_exception=_openai_handler_error(error_type, proxy_headers), + original_exception=_openai_handler_error(error_type, _PROXY_HEADERS, status_code=status_code), custom_llm_provider="litellm_proxy", completion_kwargs={}, extra_kwargs={}, @@ -1460,7 +1462,30 @@ def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str): assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" assert exc_info.value.body["type"] == error_type - assert proxy_headers.items() <= exc_info.value.headers.items() + assert exc_info.value.headers == _PROXY_HEADERS + + +@pytest.mark.parametrize( + "relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError] +) +def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_class: type[litellm.BadRequestError]): + """A proxy relaying a provider's own litellm error names the class in the message, which + re-raises that class on the SDK side before the generic 400 mapping runs; it must carry the + body and the proxy headers the same way the generic mapping now does.""" + message = f"litellm.{relayed_class.__name__}: {_GUARDRAIL_BLOCK_ERROR['message']}" + + with pytest.raises(relayed_class) as exc_info: + exception_type( + model="claude-haiku-4-5", + original_exception=_openai_handler_error("None", _PROXY_HEADERS, message=message), + custom_llm_provider="litellm_proxy", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert type(exc_info.value) is relayed_class + assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" + assert exc_info.value.headers == _PROXY_HEADERS def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index db9fe3a2253..90f1da84a61 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -145,7 +145,7 @@ def test_a_status_carried_by_an_exception_drives_the_type_it_reports(): assert openai_error_type(exc, error_status_code(exc, 400)) == "permission_error" -def test_the_stringified_none_an_older_upstream_proxy_sent_is_treated_as_absent(): +def test_a_stringified_none_type_or_param_is_treated_as_absent(): """A proxy fronting a proxy older than 1.102 receives {"type": "None", "param": "None"} on the wire; the SDK now keeps that body on the mapped exception, and re-emitting the literal is the exact bug this module exists to stop.""" From 37447c98f77116a54e3c902a64a64fae618c7996 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:28:43 -0700 Subject: [PATCH 038/107] fix(guardrails): reject per-message texts that cannot land on a string input or a Messages request --- .../chat/guardrail_translation/handler.py | 3 ++ .../base_llm/guardrail_translation/utils.py | 6 +++ .../chat/guardrail_translation/handler.py | 5 +- .../guardrail_translation/handler.py | 7 +-- .../test_anthropic_guardrail_handler.py | 52 +++++++++++++++++++ ...test_openai_responses_guardrail_handler.py | 35 +++++++++++++ 6 files changed, 102 insertions(+), 6 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50345d70d..656e9978eff 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -44,6 +44,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( scoped_structured_message_indices, stream_item_field, stream_item_fingerprint, + unappliable_request_rewrite, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -570,6 +571,8 @@ class AnthropicMessagesHandler(BaseTranslation): preserve_system_messages=has_midturn_system_message, ) else: + if guardrailed_texts and len(guardrailed_texts) != len(scanned): + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( messages=messages, diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 34d648cf184..a80e20c5404 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -411,3 +411,9 @@ def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts) rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped + + +def unappliable_request_rewrite(guardrail_name: str | None) -> Exception: + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + return UnappliableRequestRewrite(guardrail_name or "unknown") diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 56fda636e9a..ee68f8f6546 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -42,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_field, stream_item_fingerprint, stream_item_items, + unappliable_request_rewrite, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -197,9 +198,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Step 3: Map guardrail responses back to original message structure if guardrailed_texts and texts_to_check: if len(guardrailed_texts) != len(text_task_mappings): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite - - raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) await self._apply_guardrail_responses_to_input_texts( messages=messages, responses=guardrailed_texts, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 2fe11d9f7bd..27ff55f120c 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -56,6 +56,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_field, stream_item_fingerprint, stream_item_items, + unappliable_request_rewrite, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( @@ -495,13 +496,13 @@ class OpenAIResponsesHandler(BaseTranslation): data["instructions"] = written_back.instructions # rebind-ok: data is an out-param elif isinstance(input_data, str): guardrailed_texts: Final = guardrailed_inputs.get("texts") or () + if len(guardrailed_texts) > 1: + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param else: rewritten_texts: Final = guardrailed_inputs.get("texts") or () if len(rewritten_texts) != len(extracted.task_mappings): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite - - raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") + raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) await self._apply_guardrail_responses_to_input( messages=input_data, responses=rewritten_texts, diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 9fe56f4dc65..d6fd30638cc 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2272,6 +2272,58 @@ class TestAnthropicMessagesHandlerStreamingScanKey: assert ended_key != open_key +class PerRowTextGuardrail(CustomGuardrail): + """Answers one redacted text per chat row it was shown, the way a guardrail + that scans per message does, and hands back only texts.""" + + def __init__(self): + super().__init__(guardrail_name="per-row-redactor") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + rows = inputs.get("structured_messages") or [] + return {**inputs, "texts": [str(row.get("content")).replace("123-45-6789", "") for row in rows]} + + +class TestPerMessageTextWriteBack: + """Texts that no longer pair one-to-one with what the handler extracted must be + rejected by name instead of sliding onto the wrong messages.""" + + @pytest.mark.asyncio + async def test_one_text_per_row_over_a_system_prompt_is_rejected_by_name(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + data = { + "model": "claude-sonnet-4-5", + "system": "Reply with exactly the SSN you were given.", + "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}], + } + original = json.loads(json.dumps(data)) + + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail()) + + assert excinfo.value.guardrail_name == "per-row-redactor" + assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched" + assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched" + + @pytest.mark.asyncio + async def test_one_text_per_row_without_a_system_prompt_is_applied(self): + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}], + } + + await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail()) + + assert data["messages"] == [{"role": "user", "content": "My SSN is ."}] + + class TestAnthropicMessagesHandlerPostCallHookResponse: def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self): from litellm.types.utils import Choices, Message, ModelResponse, Usage 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 394f134e99c..ac719da169c 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 @@ -2392,6 +2392,14 @@ def _tool_replay_request() -> dict: } +def _string_input_request() -> dict: + return { + "model": "gpt-5.6", + "instructions": "Never repeat the SSN " + SSN + " back.", + "input": "My SSN is " + SSN + ".", + } + + class TestPerMessageRewriteWriteBack: """A guardrail that rewrites per chat row hands the rows back as structured_messages, and the handler lands them on the instructions and the @@ -2432,6 +2440,33 @@ class TestPerMessageRewriteWriteBack: assert data["input"] == original["input"] assert data["instructions"] == original["instructions"] + @pytest.mark.asyncio + async def test_structured_rows_land_on_instructions_and_string_input(self): + guardrail = _per_message_redactor() + data = _string_input_request() + + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)): + result = await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back." + assert [_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]] + + @pytest.mark.asyncio + async def test_texts_only_per_message_answer_over_a_string_input_is_rejected_by_name(self): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + guardrail = _per_message_redactor() + data = _string_input_request() + original = copy.deepcopy(data) + + with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)): + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await OpenAIResponsesHandler().process_input_messages(data, guardrail) + + assert excinfo.value.guardrail_name == "per-message-redactor" + assert data["input"] == original["input"] + assert data["instructions"] == original["instructions"] + class TestProvenancePatching: """The O(n) provenance pass must keep patching rewritten rows in place for the From dd173a0b1b7099255c51812edb327dffddb121fc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:43:35 -0700 Subject: [PATCH 039/107] fix(guardrails): count the PANW latest-user scan over the hoisted system prompt --- .../panw_prisma_airs/panw_prisma_airs.py | 11 +--- .../guardrail_hooks/test_panw_prisma_airs.py | 55 ++++++------------- 2 files changed, 20 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 3bc0dfabefc..9002e2aea07 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -1600,8 +1600,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): Args: texts: Flattened text entries from the framework. - messages: Original request messages (request_data["messages"]), - NOT structured_messages (which may have injected system content). + messages: The structured messages the framework flattened into ``texts``, + hoisted top-level system prompt included, so positions line up. Returns a set of scannable indices, or None on count mismatch or no user/developer message (safety fallback to existing role-filter behavior). @@ -1788,15 +1788,10 @@ class PanwPrismaAirsHandler(CustomGuardrail): structured_messages: Final = inputs.get("structured_messages") if structured_messages: # For Anthropic /v1/messages: default to latest-user-only scanning. - # Uses request_data["messages"] (original format), NOT structured_messages - # (which has injected system content from adapter translation). if self._use_latest_user_only(request_data, logging_obj): - original_messages: Final = request_data.get("messages") - if original_messages: - scannable_indices = self._get_latest_user_text_indices(texts, original_messages) + scannable_indices = self._get_latest_user_text_indices(texts, structured_messages) # Fall through to existing role filtering if: # - not Anthropic, OR flag explicitly False, OR - # - no original messages, OR # - latest-user extraction returned None (no user / count mismatch) if scannable_indices is None: scannable_indices = self._get_scannable_text_indices(texts, structured_messages) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 3d7c6e06d94..f25727ebd9a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -4620,46 +4620,27 @@ class TestPanwAirsLatestRoleMessageOnly: @pytest.mark.asyncio async def test_anthropic_system_plus_multiturn_no_fallback(self): - """Anthropic with top-level system + multi-turn messages[] - — latest-user works, no scan-all fallback. + """Anthropic with a top-level system prompt and multi-turn messages[] + scans only the latest user turn, with no scan-all fallback. - Key scenario: Anthropic top-level `system` field causes - structured_messages to have an injected system entry, but - request_data["messages"] does NOT include it. + The Anthropic handler hoists the top-level `system` field into both + `texts` and `structured_messages`, so the latest-user walk has to + count the same entries the framework flattened. """ - handler = PanwPrismaAirsHandler( - guardrail_name="test_panw_airs", - api_key="test_api_key", - profile_name="test_profile", - default_on=True, + from litellm.llms.anthropic.chat.guardrail_translation.handler import ( + AnthropicMessagesHandler, ) - # Original Anthropic messages (no system in messages array) - original_messages = [ - {"role": "user", "content": "First user turn"}, - {"role": "assistant", "content": "First assistant turn"}, - {"role": "user", "content": "Latest user turn"}, - ] - - # texts extracted from original_messages (3 text entries) - texts = ["First user turn", "First assistant turn", "Latest user turn"] - - # structured_messages has an INJECTED system message from translation - structured_messages = [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "First user turn"}, - {"role": "assistant", "content": "First assistant turn"}, - {"role": "user", "content": "Latest user turn"}, - ] - - inputs: GenericGuardrailAPIInputs = { - "texts": texts, - "structured_messages": structured_messages, - } + handler = make_handler() request_data = { "litellm_call_id": "test-call-id", "model": "anthropic/claude-sonnet-4-20250514", - "messages": original_messages, + "system": "You are a helpful assistant.", + "messages": [ + {"role": "user", "content": "First user turn"}, + {"role": "assistant", "content": "First assistant turn"}, + {"role": "user", "content": "Latest user turn"}, + ], "proxy_server_request": { "url": "http://localhost:4000/v1/messages", }, @@ -4670,13 +4651,11 @@ class TestPanwAirsLatestRoleMessageOnly: ) as mock_api: mock_api.return_value = {"action": "allow", "category": "benign"} - await handler.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="request", + await AnthropicMessagesHandler().process_input_messages( + data=request_data, + guardrail_to_apply=handler, ) - # Should scan ONLY the latest user message, not fall back to scan-all assert mock_api.call_count == 1 assert mock_api.call_args.kwargs["content"] == "Latest user turn" From 6264bd84bf1e0b33865028acdf53f5d64adb1e98 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:43:36 -0700 Subject: [PATCH 040/107] chore(ui): regenerate dashboard API types --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From 7ea19eccc7a6ebe4f6c0bab35073f02b17407fe5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:02:03 -0700 Subject: [PATCH 041/107] fix(responses): keep namespace custom tools through guardrail merges and Mantle params identity --- .../responses/transformation.py | 9 +++-- .../guardrail_translation/tool_merge.py | 31 ++++++++-------- .../custom_tools.py | 35 +++++++++++++------ .../transformation.py | 3 +- ...bedrock_mantle_responses_transformation.py | 6 ++++ ...t_openai_responses_guardrail_tool_merge.py | 32 +++++++++++++++-- 6 files changed, 85 insertions(+), 31 deletions(-) diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 95375399033..57590601a3c 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -237,10 +237,15 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) hoisted: Final = hoist_additional_tools(input, params.get("tools")) normalized_input: Final = self._normalize_codex_input_items(hoisted.input) + request_params: Final = ( + self._params_with_hoisted_tools(params, hoisted) + if hoisted.hoisted + else response_api_optional_request_params + ) return super().transform_responses_api_request( model=model, input=normalized_input, - response_api_optional_request_params=self._params_with_hoisted_tools(params, hoisted), + response_api_optional_request_params=request_params, litellm_params=litellm_params, headers=headers, ) @@ -249,8 +254,6 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI def _params_with_hoisted_tools( cls, params: Mapping[str, object], hoisted: HoistedAdditionalTools ) -> dict[str, object]: - if not hoisted.hoisted: - return dict(params) supported_tools: Final = cls._filter_unsupported_tools(list(hoisted.tools)) if supported_tools: return {**params, "tools": supported_tools} diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index b596adfad6f..0326e9b2bfd 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -6,8 +6,10 @@ from typing import Final, TypeAlias from pydantic import BaseModel, TypeAdapter, ValidationError from litellm._logging import verbose_logger +from litellm.responses.litellm_completion_transformation.custom_tools import custom_tool_grammar_suffix from litellm.responses.litellm_completion_transformation.transformation import ( NAMESPACE_DESCRIPTION_SEPARATOR, + NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS, LiteLLMCompletionResponsesConfig, ) @@ -34,8 +36,8 @@ def _validated_tools(values: Iterable[object]) -> tuple[Tool, ...]: return tuple(tool for tool in validated if tool is not None) -def _is_function(tool: Tool) -> bool: - return tool.get("type") == "function" +def _has_chat_tool(member: Tool) -> bool: + return member.get("type") in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS def _chat_tool_key(tool: Tool) -> str: @@ -67,18 +69,19 @@ def _function_fields(tool: Tool) -> Tool: return function if function is not None else MappingProxyType({}) -def _without_namespace_prefix(key: str, value: object, prefix: str) -> object: - if key != "description" or not isinstance(value, str) or not value.startswith(prefix): +def _member_description(key: str, value: object, prefix: str, suffix: str) -> object: + if key != "description" or not isinstance(value, str): return value - return value[len(prefix) :] + return value.removeprefix(prefix).removesuffix(suffix) def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool: flattened_function: Final = _function_fields(flattened) prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else "" + suffix: Final = custom_tool_grammar_suffix(member.get("format")) if member.get("type") == "custom" else "" changed_function: Final = MappingProxyType( { - key: _without_namespace_prefix(key, value, prefix) + key: _member_description(key, value, prefix, suffix) for key, value in _function_fields(guardrailed).items() if flattened_function.get(key) != value } @@ -93,8 +96,8 @@ def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_ return {**member, **changed_extras, **changed_function} # mutable-ok: json.dumps rejects MappingProxyType -def _rebuilt_function_members( - function_members: Sequence[Tool], +def _rebuilt_flattened_members( + flattened_members: Sequence[Tool], flattened_group: Sequence[Tool], group_keys: Sequence[IndexedKey], guardrailed_by_key: Mapping[IndexedKey, Tool], @@ -106,7 +109,7 @@ def _rebuilt_function_members( else member if guardrailed_by_key[key] == flattened else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description) - for member, flattened, key in zip(function_members, flattened_group, group_keys) + for member, flattened, key in zip(flattened_members, flattened_group, group_keys) ) @@ -118,9 +121,9 @@ def _rebuilt_namespace( guardrailed_by_key: Mapping[IndexedKey, Tool], ) -> tuple[Tool, ...]: namespace_description: Final = str(original.get("description") or "") - rebuilt_functions: Final = iter( - _rebuilt_function_members( - tuple(member for member in members if _is_function(member)), + rebuilt_flattened: Final = iter( + _rebuilt_flattened_members( + tuple(member for member in members if _has_chat_tool(member)), flattened_group, group_keys, guardrailed_by_key, @@ -129,7 +132,7 @@ def _rebuilt_namespace( ) rebuilt_members: Final = tuple( rebuilt - for rebuilt in (next(rebuilt_functions) if _is_function(member) else member for member in members) + for rebuilt in (next(rebuilt_flattened) if _has_chat_tool(member) else member for member in members) if rebuilt is not None ) if not rebuilt_members: @@ -149,7 +152,7 @@ def _merged_original( if guardrailed_group == tuple(flattened_group): return (original,) members: Final = _namespace_members(original) if original.get("type") == "namespace" else () - if members and sum(map(_is_function, members)) == len(flattened_group): + if members and sum(map(_has_chat_tool, members)) == len(flattened_group): return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key) if not guardrailed_group: return () diff --git a/litellm/responses/litellm_completion_transformation/custom_tools.py b/litellm/responses/litellm_completion_transformation/custom_tools.py index 038964055c3..7888a07e248 100644 --- a/litellm/responses/litellm_completion_transformation/custom_tools.py +++ b/litellm/responses/litellm_completion_transformation/custom_tools.py @@ -45,21 +45,32 @@ class _ToolNameFields(BaseModel): tools: tuple[object, ...] = () -def _custom_tool_names_of(tool: object) -> tuple[str, ...]: +def _tool_name_fields_of(tool: object) -> _ToolNameFields | None: try: - parsed: Final = _ToolNameFields.model_validate(tool) + return _ToolNameFields.model_validate(tool) except ValidationError: + return None + + +def _custom_tool_name_of(tool: object) -> str | None: + parsed: Final = _tool_name_fields_of(tool) + if parsed is None or parsed.type != "custom" or not parsed.name: + return None + return parsed.name + + +def _nested_tools_of(tool: object) -> tuple[object, ...]: + parsed: Final = _tool_name_fields_of(tool) + if parsed is None or parsed.type != "namespace": return () - if parsed.type == "custom": - return (parsed.name,) if parsed.name else () - if parsed.type != "namespace": - return () - return tuple(name for nested_tool in parsed.tools for name in _custom_tool_names_of(nested_tool)) + return parsed.tools def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]: - """Extract names of tools defined as ``type: "custom"``, at the top level or inside a ``namespace`` tool.""" - return {name for tool in tools or () for name in _custom_tool_names_of(tool)} + """Extract names of ``type: "custom"`` tools, at the top level or one level inside a ``namespace`` tool.""" + top_level: Final = tuple(tools or ()) + nested: Final = tuple(nested_tool for tool in top_level for nested_tool in _nested_tools_of(tool)) + return {name for tool in (*top_level, *nested) if (name := _custom_tool_name_of(tool)) is not None} def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool: @@ -155,7 +166,7 @@ def validated_allowed_callers(value: object) -> list[str] | None: raise ValueError("allowed_callers must be a list of strings") from exc -def _grammar_suffix(fmt: object) -> str: +def custom_tool_grammar_suffix(fmt: object) -> str: try: parsed: Final = _CustomToolFormat.model_validate(fmt) except ValidationError: @@ -179,7 +190,9 @@ def convert_custom_tool_to_function_tool(tool: Mapping[str, object]) -> ChatComp raw_name: Final = tool.get("name") name: Final = raw_name if isinstance(raw_name, str) else "" raw_description: Final = tool.get("description") - description = (raw_description if isinstance(raw_description, str) else "") + _grammar_suffix(tool.get("format")) + description: Final = (raw_description if isinstance(raw_description, str) else "") + custom_tool_grammar_suffix( + tool.get("format") + ) allowed_callers: Final = validated_allowed_callers(tool.get("allowed_callers")) function_chunk: Final = ChatCompletionToolParamFunctionChunk( name=name, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 27756b405ec..d27fc855be6 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -109,6 +109,7 @@ NamespaceTool: TypeAlias = Mapping[str, object] ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" +NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"}) @dataclass(frozen=True, slots=True) @@ -1891,7 +1892,7 @@ class LiteLLMCompletionResponsesConfig: nested: bool, ) -> ChatCompletionToolParam | None: tool_type: Final = namespace_tool.get("type") - if nested and tool_type not in ("function", "custom"): + if nested and tool_type not in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: return None raw_description: Final = str(namespace_tool.get("description") or "") 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 40566261c84..a7aefa714aa 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 @@ -773,6 +773,12 @@ class TestBedrockMantleCodexAdditionalTools: assert body["input"] == codex_agentic_items assert "tools" not in body + def test_input_without_additional_tools_sanitizes_tools_on_the_caller_params_object(self): + params = {"tools": [{"type": "function", "name": "wait", "parameters": '{"type": "object"}'}]} + body = self._transform(input=[self._USER_MESSAGE], params=params) + assert body["tools"][0]["parameters"] == {"type": "object"} + assert params["tools"][0]["parameters"] == {"type": "object"} + def test_malformed_additional_tools_item_without_tools_list_is_stripped(self): body = self._transform( input=[ diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py index 9c236d81f51..a7f65545eb2 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -133,7 +133,20 @@ def test_namespace_keeps_a_non_function_member_when_a_function_member_is_edited( assert merged[0]["tools"][1] == custom_member -def test_namespace_keeps_its_non_function_members_when_every_function_member_is_dropped(): +def test_namespace_keeps_its_custom_member_when_every_function_member_is_dropped(): + custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}} + original = [ + {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read"), custom_member]}, + _function("a"), + ] + groups = _groups(original) + + merged = merge_guardrailed_tools(original, groups, [groups[0][1], groups[1][0]]) + + assert list(merged) == [{"type": "namespace", "name": "ns", "description": "NS", "tools": [custom_member]}, _function("a")] + + +def test_namespace_custom_member_is_dropped_when_the_guardrail_drops_its_chat_form(): custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}} original = [ {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read"), custom_member]}, @@ -143,7 +156,22 @@ def test_namespace_keeps_its_non_function_members_when_every_function_member_is_ merged = merge_guardrailed_tools(original, groups, [groups[1][0]]) - assert list(merged) == [{"type": "namespace", "name": "ns", "description": "NS", "tools": [custom_member]}, _function("a")] + assert list(merged) == [_function("a")] + + +def test_custom_member_description_edit_lands_without_the_namespace_prefix_or_grammar_block(): + grammar = {"type": "grammar", "syntax": "lark", "definition": "start: X"} + custom_member = {"type": "custom", "name": "exec", "description": "Run a command", "format": grammar} + original = [{"type": "namespace", "name": "shell", "description": "Shell", "tools": [custom_member]}] + groups = _groups(original) + assert groups[0][0]["function"]["description"] == "Shell\n\nRun a command\n\nFormat:\n```lark\nstart: X\n```" + edited = copy.deepcopy(_flat(groups)) + edited[0]["function"]["description"] = "Shell\n\nRun a command (guarded)\n\nFormat:\n```lark\nstart: X\n```" + + merged = merge_guardrailed_tools(original, groups, edited) + + guarded_member = {**custom_member, "description": "Run a command (guarded)"} + assert list(merged) == [{"type": "namespace", "name": "shell", "description": "Shell", "tools": [guarded_member]}] def test_member_extras_edited_by_the_guardrail_land_on_that_member(): From 2923c4ac5520642bc3c3e728f9b2974850a43f7e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:02:15 -0700 Subject: [PATCH 042/107] fix(guardrails): read the rewrite from texts when a guardrail echoes every row back unchanged --- .../generic_guardrail_api.py | 25 ++++++++------ .../test_generic_guardrail_api.py | 34 +++++++++++++++++++ 2 files changed, 49 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 16159d32a7f..3d1a173635e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -150,16 +150,20 @@ def _extract_inbound_headers( return None -def _rows_with_unchanged_originals( +def _structured_rows_to_write_back( original_rows: Sequence[AllMessageValues] | None, shown_rows: Sequence[AllMessageValues] | None, returned_rows: Sequence[AllMessageValues], -) -> tuple[AllMessageValues, ...]: +) -> tuple[AllMessageValues, ...] | None: """The request model drops row keys its message types do not declare, so a - row the server echoes back verbatim is restored to the original row object; - only rows the server actually changed reach the endpoint write-back.""" + row the server echoes back verbatim is restored to the original row object. + A server that echoes every row back unchanged has not rewritten anything + per row, so its answer is read from texts, as it was before rows could be + returned at all.""" if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows): return tuple(returned_rows) + if all(returned == shown for shown, returned in zip(shown_rows, returned_rows)): + return None return tuple( original if returned == shown else returned for original, shown, returned in zip(original_rows, shown_rows, returned_rows) @@ -354,12 +358,13 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools - if guardrail_response.structured_messages: - return_inputs["structured_messages"] = list( # mutable-ok: guardrail inputs take a list - _rows_with_unchanged_originals( - structured_messages, shown_messages, guardrail_response.structured_messages - ) - ) + rows_to_write_back: Final = ( + _structured_rows_to_write_back(structured_messages, shown_messages, guardrail_response.structured_messages) + if guardrail_response.structured_messages + else None + ) + if rows_to_write_back is not None: + return_inputs["structured_messages"] = list(rows_to_write_back) # mutable-ok: guardrail inputs take a list if guardrail_response.stream_holdback_chars is not None: return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars return return_inputs diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index cc9942e0e40..a5e79f84ef1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -659,6 +659,40 @@ class TestStructuredMessagesInResponse: assert returned_rows[1] is tool_call_row assert returned_rows[2] == {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'} + @pytest.mark.asyncio + async def test_rows_all_echoed_back_as_shown_leave_the_rewrite_to_texts( + self, generic_guardrail, mock_request_data_input + ): + """A server written against the texts contract that echoes the request rows back + untouched while rewriting texts still gets its texts rewrite applied.""" + original_rows = [ + {"role": "system", "content": "Never repeat an SSN."}, + {"role": "user", "content": "Look up 123-45-6789 for me."}, + ] + + def echo_rows_and_rewrite_texts(url, json, headers): + answer = MagicMock() + answer.json.return_value = { + "action": "NONE", + "texts": [text.replace("123-45-6789", "[REDACTED]") for text in json["texts"]], + "structured_messages": json["structured_messages"], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object(generic_guardrail.async_handler, "post", side_effect=echo_rows_and_rewrite_texts): + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={ + "texts": ["Never repeat an SSN.", "Look up 123-45-6789 for me."], + "structured_messages": original_rows, + }, + request_data=mock_request_data_input, + input_type="request", + ) + + assert "structured_messages" not in guardrailed_inputs + assert guardrailed_inputs["texts"] == ["Never repeat an SSN.", "Look up [REDACTED] for me."] + @pytest.mark.asyncio @pytest.mark.parametrize( "structured_messages", From ff5b59b17336163ac251dc89700b1af406982947 Mon Sep 17 00:00:00 2001 From: ryan Date: Mon, 14 Sep 2026 01:17:19 +0000 Subject: [PATCH 043/107] feat(proxy): add POST /user/bulk_new for batched user and team membership creation Creates up to 500 internal users in one request with set-based validation, a single create_many for user rows, and one locked write per referenced team. Rows fail independently, keys are opt-in per row via auto_create_key, and send_invite_email is rejected for the batch. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 1 + litellm/proxy/auth/route_checks.py | 2 + .../internal_user_endpoints.py | 68 ++ .../key_management_endpoints.py | 69 +- .../management_helpers/bulk_user_creation.py | 820 ++++++++++++++++++ .../internal_user_endpoints.py | 40 +- .../test_bulk_user_creation.py | 347 ++++++++ 7 files changed, 1321 insertions(+), 26 deletions(-) create mode 100644 litellm/proxy/management_helpers/bulk_user_creation.py create mode 100644 tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b315a2beac9..ed182460644 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -656,6 +656,7 @@ class LiteLLMRoutes(enum.Enum): [ # user "/user/new", + "/user/bulk_new", "/user/update", "/user/bulk_update", "/user/delete", diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 953e3cf3e88..051ab13c058 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -24,6 +24,7 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset( [ # user "/user/new", + "/user/bulk_new", "/user/delete", "/user/bulk_update", # team @@ -755,6 +756,7 @@ class RouteChecks: _ADMIN_VIEWER_BLOCKED_WRITE_ROUTES = frozenset( [ "/user/new", + "/user/bulk_new", "/user/delete", "/user/bulk_update", "/team/new", diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index e3efda507f6..2624d7e373a 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -5,6 +5,7 @@ Internal User Management Endpoints These are members of a Team on LiteLLM /user/new +/user/bulk_new /user/update /user/bulk_update /user/delete @@ -77,6 +78,8 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( + BulkNewUserRequest, + BulkNewUserResponse, BulkUpdateUserRequest, BulkUpdateUserResponse, UserListResponse, @@ -637,6 +640,71 @@ async def new_user( raise handle_exception_on_proxy(e) +@router.post( + "/user/bulk_new", + tags=["Internal User management"], + dependencies=[Depends(user_api_key_auth)], + response_model=BulkNewUserResponse, +) +@management_endpoint_wrapper +async def bulk_new_user( + data: BulkNewUserRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection +) -> BulkNewUserResponse: + """ + Create up to 500 internal users in one request, optionally adding each one to teams. + + Every entry in `users` takes the same fields as `/user/new`, with two differences: `auto_create_key` + defaults to `false` (opt in per user to also get a virtual key back) and `send_invite_email` is not + supported. Rows are validated together (duplicate ids or emails, unknown teams, roles the caller may not + grant), inserted in one statement, and each referenced team is written once for all of its new members. + + Rows fail independently: a bad row is reported in `results` with `success: false` and an `error`, and the + other rows still get created. A user that was created but could not be added to one of its teams is + reported with `success: true`, `teams` listing where they did land, and `error` naming the failed team. + The whole request is rejected with 403 only if creating the valid rows would exceed the license seat limit. + + Usage Example + + ```shell + curl -X POST "http://localhost:4000/user/bulk_new" \\ + -H "Content-Type: application/json" \\ + -H "Authorization: Bearer sk-1234" \\ + -d '{ + "users": [ + {"user_email": "a@example.com", "user_role": "internal_user", "teams": ["team-1"]}, + {"user_email": "b@example.com", "user_role": "internal_user", "auto_create_key": true} + ] + }' + ``` + + Returns `results` (one entry per input row, in order, with `user_id`, `user_email`, `success`, `teams`, + `key`, `error`), `total_requested`, `successful_creations` and `failed_creations`. + """ + from litellm.proxy.management_helpers.bulk_user_creation import bulk_create_users + from litellm.proxy.proxy_server import ( + _license_check, # pyright: ignore[reportPrivateUsage] # same proxy license singleton /user/new reads + litellm_proxy_admin_name, + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + raise HTTPException(status_code=400, detail=CommonProxyErrors.db_not_connected_error.value) + try: + return await bulk_create_users( + users=data.users, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + license_check=_license_check, + litellm_proxy_admin_name=litellm_proxy_admin_name, + user_api_key_cache=user_api_key_cache, + ) + except Exception as e: # noqa: BLE001 # normalize every failure to the proxy exception contract + verbose_proxy_logger.exception("/user/bulk_new: Exception occured - %s", e) + raise handle_exception_on_proxy(e) + + @router.get( "/user/available_roles", tags=["Internal User management"], diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 95ccb7bbe0b..11dd7ce1fa5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4082,6 +4082,40 @@ def _check_model_access_group(models: list[str] | None, llm_router: Router | Non return True +_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) + + +def metadata_json_with_limits( + metadata: Mapping[str, object] | None, + *, + model_rpm_limit: Mapping[str, object] | None, + model_tpm_limit: Mapping[str, object] | None, + mcp_rpm_limit: Mapping[str, int] | None, + tag_rpm_limit: Mapping[str, int] | None, + guardrails: Sequence[str] | None, + policies: Sequence[str] | None, + prompts: Sequence[str] | None, +) -> str: + """Serialize the stored metadata blob with the per-model, MCP, tag, guardrail, policy and prompt settings folded in.""" + limits: Final = tuple( + (name, value) + for name, value in ( + ("model_rpm_limit", model_rpm_limit), + ("model_tpm_limit", model_tpm_limit), + ("mcp_rpm_limit", mcp_rpm_limit), + ("tag_rpm_limit", tag_rpm_limit), + ("guardrails", guardrails), + ("policies", policies), + ("prompts", prompts), + ) + if value is not None + ) + if metadata is None and not limits: + return json.dumps(None) + merged: Final = {**(metadata or _NO_METADATA), **dict(limits)} # mutable-ok: encrypt_callback_vars takes a dict + return json.dumps(encrypt_callback_vars(merged)) + + async def generate_key_helper_fn( request_type: Literal["user", "key"], # identifies if this request is from /user/new or /key/generate duration: str | None = None, @@ -4184,31 +4218,16 @@ async def generate_key_helper_fn( permissions_json: Final = json.dumps(permissions) router_settings_json: Final = safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) - # Add model_rpm_limit and model_tpm_limit to metadata - if model_rpm_limit is not None: - metadata = metadata or {} - metadata["model_rpm_limit"] = model_rpm_limit - if model_tpm_limit is not None: - metadata = metadata or {} - metadata["model_tpm_limit"] = model_tpm_limit - if mcp_rpm_limit is not None: - metadata = metadata or {} - metadata["mcp_rpm_limit"] = mcp_rpm_limit - if tag_rpm_limit is not None: - metadata = metadata or {} - metadata["tag_rpm_limit"] = tag_rpm_limit - if guardrails is not None: - metadata = metadata or {} - metadata["guardrails"] = guardrails - if policies is not None: - metadata = metadata or {} - metadata["policies"] = policies - if prompts is not None: - metadata = metadata or {} - metadata["prompts"] = prompts - - metadata = encrypt_callback_vars(metadata) - metadata_json: Final = json.dumps(metadata) + metadata_json: Final = metadata_json_with_limits( + metadata, + model_rpm_limit=model_rpm_limit, + model_tpm_limit=model_tpm_limit, + mcp_rpm_limit=mcp_rpm_limit, + tag_rpm_limit=tag_rpm_limit, + guardrails=guardrails, + policies=policies, + prompts=prompts, + ) validate_model_max_budget(model_max_budget) model_max_budget_json: Final = json.dumps(model_max_budget) budget_fallbacks_json: Final = json.dumps(budget_fallbacks or {}) diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py new file mode 100644 index 00000000000..301629e6c2a --- /dev/null +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -0,0 +1,820 @@ +"""Batched internal user creation behind `/user/bulk_new`. + +The batch is validated with set queries, user rows land in one `create_many`, and every +referenced team is written once under its advisory lock instead of once per user. +""" + +import asyncio +import json +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, TypeVar + +from fastapi import HTTPException, Request +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import ( + LiteLLM_TeamTable, + LitellmUserRoles, + Member, + NewUserRequestTeam, + OrganizationMemberAddRequest, + OrgMember, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import invalidate_team_member_spend_state +from litellm.proxy.auth.litellm_license import LicenseCheck +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks +from litellm.proxy.management_endpoints.common_utils import ( + _is_user_org_admin_for_team, # pyright: ignore[reportPrivateUsage] # same team-admin check /user/new uses + _is_user_team_admin, # pyright: ignore[reportPrivateUsage] # same team-admin check /user/new uses + validate_budget_duration, +) +from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_internal_new_user_params, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # /user/new defaults; result validated below + check_if_default_team_set, +) +from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_permissions_caller_permission, # pyright: ignore[reportPrivateUsage] # same permission check /user/new uses + generate_key_helper_fn, # pyright: ignore[reportUnknownVariableType] # legacy untyped helper; result validated by _KEY_RESPONSE + metadata_json_with_limits, +) +from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add +from litellm.proxy.management_helpers.access_group_team_sync import TEAM_ADVISORY_LOCK_SQL +from litellm.proxy.management_helpers.object_permission_utils import ( + _set_object_permission, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared with /user/new; result validated below +) +from litellm.proxy.management_helpers.utils import ( + _resolve_member_budget_id, # pyright: ignore[reportPrivateUsage] # shared with /team/member_add +) +from litellm.proxy.utils import PrismaClient +from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( + BulkNewUserItem, + BulkNewUserResponse, + UserCreateResult, +) + +if TYPE_CHECKING: + from prisma import Prisma + from prisma import models as prisma_models + + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + +BULK_NEW_USER_CONCURRENCY: Final = 10 + +TeamRole: TypeAlias = Literal["user", "admin"] +KeyGenerator: TypeAlias = Callable[..., Awaitable[object]] +_T: Final = TypeVar("_T") + + +@dataclass(frozen=True, slots=True) +class _RowFailure: + index: int + user_id: str | None + user_email: str | None + error: str + + +@dataclass(frozen=True, slots=True) +class _PendingUser: + index: int + request: BulkNewUserItem + user_id: str + teams: tuple[NewUserRequestTeam, ...] + + +class _UserRow(BaseModel): + """The `/user/new` body after defaults and object permission were applied.""" + + model_config = ConfigDict(extra="ignore") + + user_id: str + user_email: str | None = None + user_alias: str | None = None + user_role: str | None = None + team_id: str | None = None + max_budget: float | None = None + spend: float | None = 0.0 + models: tuple[str, ...] | None = None + metadata: Mapping[str, object] | None = None + max_parallel_requests: int | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + budget_duration: str | None = None + allowed_cache_controls: tuple[str, ...] | None = None + sso_user_id: str | None = None + object_permission_id: str | None = None + model_max_budget: Mapping[str, object] | None = None + model_rpm_limit: Mapping[str, object] | None = None + model_tpm_limit: Mapping[str, object] | None = None + mcp_rpm_limit: Mapping[str, int] | None = None + tag_rpm_limit: Mapping[str, int] | None = None + guardrails: tuple[str, ...] | None = None + policies: tuple[str, ...] | None = None + prompts: tuple[str, ...] | None = None + duration: str | None = None + key_alias: str | None = None + organizations: tuple[str, ...] | None = None + + +_USER_ROW: Final = TypeAdapter(_UserRow) + + +@dataclass(frozen=True, slots=True) +class _PreparedUser: + pending: _PendingUser + row: _UserRow + + +@dataclass(frozen=True, slots=True) +class _TeamAssignment: + user_id: str + user_email: str | None + role: TeamRole + max_budget_in_team: float | None + + +@dataclass(frozen=True, slots=True) +class _TeamWrite: + """Outcome of one locked roster write. `failed` maps user ids to the reason they were not added.""" + + team_id: str + after: tuple[Member, ...] + added: frozenset[str] + failed: Mapping[str, str] + + +@dataclass(frozen=True, slots=True) +class _CreatedUser: + prepared: _PreparedUser + teams: tuple[str, ...] + key: str | None + errors: tuple[str, ...] + + +_ERROR_DETAIL: Final = TypeAdapter(Mapping[str, object]) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) + + +class _KeyResponse(BaseModel): + token: str + + +_KEY_RESPONSE: Final = TypeAdapter(_KeyResponse) + + +def _error_message(exc: BaseException) -> str: + if not isinstance(exc, HTTPException): + return str(exc) + try: + detail: Final = _ERROR_DETAIL.validate_python(exc.detail) + except ValidationError: + return str(exc.detail) + return str(detail.get("error", detail)) + + +def _requested_teams(item: BulkNewUserItem) -> tuple[NewUserRequestTeam, ...]: + if item.team_id is not None: + return (NewUserRequestTeam(team_id=item.team_id),) + teams: Final = item.teams if item.teams is not None else check_if_default_team_set() + if teams is None: + return () + return tuple(team if isinstance(team, NewUserRequestTeam) else NewUserRequestTeam(team_id=team) for team in teams) + + +def _row_error(item: BulkNewUserItem, user_api_key_dict: UserAPIKeyAuth) -> str | None: + if ( + item.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN + ): + return ( + "Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). " + f"Attempted to create user with role: {item.user_role}. Your role: {user_api_key_dict.user_role}" + ) + try: + validate_budget_duration(item.budget_duration) + _check_permissions_caller_permission(data=item, user_api_key_dict=user_api_key_dict) + except Exception as exc: # noqa: BLE001 # any validation failure is reported on this row only + return _error_message(exc) + return None + + +def _normalized_email(email: str | None) -> str | None: + return email.strip().lower() if email else None + + +def _partition_rows( + users: Sequence[BulkNewUserItem], user_api_key_dict: UserAPIKeyAuth +) -> tuple[tuple[_PendingUser, ...], tuple[_RowFailure, ...]]: + """Assign ids, run the per-row checks and fail later rows that repeat an earlier row's id or email.""" + user_ids: Final = tuple(item.user_id or str(uuid.uuid4()) for item in users) + first_index_by_id: Final = MappingProxyType( + {user_id: index for index, user_id in reversed(tuple(enumerate(user_ids)))} + ) + first_index_by_email: Final = MappingProxyType( + { + email: index + for index, email in reversed(tuple(enumerate(_normalized_email(item.user_email) for item in users))) + if email is not None + } + ) + + def classify(index: int, item: BulkNewUserItem) -> _PendingUser | _RowFailure: + user_id: Final = user_ids[index] + email: Final = _normalized_email(item.user_email) + if first_index_by_id[user_id] != index: + return _RowFailure(index, user_id, item.user_email, f"Duplicate user_id in request: {user_id}") + if email is not None and first_index_by_email[email] != index: + return _RowFailure(index, user_id, item.user_email, f"Duplicate user_email in request: {item.user_email}") + error: Final = _row_error(item, user_api_key_dict) + if error is not None: + return _RowFailure(index, user_id, item.user_email, error) + return _PendingUser(index, item, user_id, _requested_teams(item)) + + outcomes: Final = tuple(classify(index, item) for index, item in enumerate(users)) + return ( + tuple(outcome for outcome in outcomes if isinstance(outcome, _PendingUser)), + tuple(outcome for outcome in outcomes if isinstance(outcome, _RowFailure)), + ) + + +def _user_table(prisma_client: PrismaClient) -> "TableActions[prisma_models.LiteLLM_UserTable]": + return UserRepository(prisma_client).table + + +async def _existing_user_conflicts( + prisma_client: PrismaClient, pending: Sequence[_PendingUser] +) -> tuple[frozenset[str], frozenset[str]]: + """Return the requested user ids and (lowercased) emails that already exist, using one query each.""" + user_ids: Final = sorted(user.user_id for user in pending) + emails: Final = sorted(frozenset(user.request.user_email for user in pending if user.request.user_email)) + if not user_ids: + return frozenset(), frozenset() + table: Final = _user_table(prisma_client) + id_filter: Final = {"user_id": {"in": user_ids}} # mutable-ok: Prisma query filters are dict-shaped + email_filter: Final = {"user_email": {"in": emails, "mode": "insensitive"}} # mutable-ok: Prisma filter + id_rows: Final = await table.find_many(where=id_filter) + email_rows: Final = await table.find_many(where=email_filter) if emails else () + return ( + frozenset(row.user_id for row in id_rows), + frozenset(lowered for row in email_rows if (lowered := _normalized_email(row.user_email)) is not None), + ) + + +async def _load_teams(prisma_client: PrismaClient, team_ids: frozenset[str]) -> Mapping[str, LiteLLM_TeamTable]: + if not team_ids: + return MappingProxyType({}) + rows: Final = await TeamRepository(prisma_client).table.find_many( + where={"team_id": {"in": sorted(team_ids)}} # mutable-ok: Prisma query filters are dict-shaped + ) + return MappingProxyType({row.team_id: LiteLLM_TeamTable.model_validate(row.model_dump()) for row in rows}) + + +async def _team_permission_error(team: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth) -> str | None: + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return None + if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team): + return None + if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team): + return None + return f"Call not allowed. User not proxy admin OR team admin. team_id={team.team_id}" + + +async def _unusable_teams( + prisma_client: PrismaClient, + pending: Sequence[_PendingUser], + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[Mapping[str, LiteLLM_TeamTable], Mapping[str, str]]: + """Load every referenced team once and explain, per team id, why rows naming it cannot proceed.""" + team_ids: Final = frozenset(team.team_id for user in pending for team in user.teams) + teams: Final = await _load_teams(prisma_client, team_ids) + permission_errors: Final = await asyncio.gather( + *(_team_permission_error(team, user_api_key_dict) for team in teams.values()) + ) + missing: Final = tuple( + (team_id, f"Team id={team_id} does not exist") for team_id in team_ids if team_id not in teams + ) + denied: Final = tuple( + (team.team_id, error) + for team, error in zip(teams.values(), permission_errors, strict=True) + if error is not None + ) + return teams, MappingProxyType({team_id: error for team_id, error in (*missing, *denied)}) + + +def _db_failure( + user: _PendingUser, + existing_ids: frozenset[str], + existing_emails: frozenset[str], + team_errors: Mapping[str, str], +) -> _RowFailure | None: + email: Final = _normalized_email(user.request.user_email) + if user.user_id in existing_ids: + return _RowFailure(user.index, user.user_id, user.request.user_email, f"User id={user.user_id} already exists") + if email is not None and email in existing_emails: + return _RowFailure( + user.index, user.user_id, user.request.user_email, f"User email={user.request.user_email} already exists" + ) + errors: Final = tuple(team_errors[team.team_id] for team in user.teams if team.team_id in team_errors) + if errors: + return _RowFailure(user.index, user.user_id, user.request.user_email, "; ".join(errors)) + return None + + +async def _prepare_user(user: _PendingUser, prisma_client: PrismaClient) -> _PreparedUser | _RowFailure: + try: + dumped: Final = user.request.model_dump(exclude={"user_id"}) # mutable-ok: pydantic IncEx takes a set + data: Final = {**dumped, "user_id": user.user_id} # mutable-ok: /user/new defaults helper mutates in place + data_json: Final = _JSON_OBJECT.validate_python(_update_internal_new_user_params(data, user.request)) + with_permission: Final = _JSON_OBJECT.validate_python( + await _set_object_permission(data_json=data_json, prisma_client=prisma_client) # pyright: ignore[reportUnknownArgumentType] # validated by the adapter + ) + return _PreparedUser(user, _USER_ROW.validate_python(with_permission)) + except Exception as exc: # noqa: BLE001 # any preparation failure is reported on this row only + verbose_proxy_logger.warning("/user/bulk_new: could not prepare user %s - %s", user.user_id, exc) + return _RowFailure(user.index, user.user_id, user.request.user_email, _error_message(exc)) + + +class _UserCreateData(TypedDict): + """One `LiteLLM_UserTable` row as `create_many` takes it; JSON columns are pre-serialized.""" + + user_id: ReadOnly[str] + user_email: ReadOnly[str | None] + user_alias: ReadOnly[str | None] + user_role: ReadOnly[str | None] + team_id: ReadOnly[str | None] + max_budget: ReadOnly[float | None] + spend: ReadOnly[float] + models: ReadOnly[tuple[str, ...]] + metadata: ReadOnly[str] + max_parallel_requests: ReadOnly[int | None] + tpm_limit: ReadOnly[int | None] + rpm_limit: ReadOnly[int | None] + budget_duration: ReadOnly[str | None] + budget_reset_at: ReadOnly[datetime | None] + allowed_cache_controls: ReadOnly[tuple[str, ...]] + sso_user_id: ReadOnly[str | None] + object_permission_id: ReadOnly[str | None] + teams: ReadOnly[tuple[str, ...]] + model_max_budget: ReadOnly[str] + + +def _user_create_payload(prepared: _PreparedUser) -> _UserCreateData: + row: Final = prepared.row + metadata_json: Final = metadata_json_with_limits( + row.metadata, + model_rpm_limit=row.model_rpm_limit, + model_tpm_limit=row.model_tpm_limit, + mcp_rpm_limit=row.mcp_rpm_limit, + tag_rpm_limit=row.tag_rpm_limit, + guardrails=row.guardrails, + policies=row.policies, + prompts=row.prompts, + ) + payload: Final[_UserCreateData] = { + "user_id": row.user_id, + "user_email": row.user_email, + "user_alias": row.user_alias, + "user_role": row.user_role, + "team_id": row.team_id, + "max_budget": row.max_budget, + "spend": row.spend or 0.0, + "models": row.models or (), + "metadata": metadata_json, + "max_parallel_requests": row.max_parallel_requests, + "tpm_limit": row.tpm_limit, + "rpm_limit": row.rpm_limit, + "budget_duration": row.budget_duration, + "budget_reset_at": get_budget_reset_time(row.budget_duration) if row.budget_duration else None, + "allowed_cache_controls": row.allowed_cache_controls or (), + "sso_user_id": row.sso_user_id, + "object_permission_id": row.object_permission_id, + "teams": tuple(team.team_id for team in prepared.pending.teams), + "model_max_budget": json.dumps(row.model_max_budget) if row.model_max_budget else "{}", + } + return payload + + +async def _bounded(limit: int, awaitables: Sequence[Awaitable[_T]]) -> tuple[_T | BaseException, ...]: + semaphore: Final = asyncio.Semaphore(limit) + + async def run(awaitable: Awaitable[_T]) -> _T: + async with semaphore: + return await awaitable + + return tuple(await asyncio.gather(*(run(awaitable) for awaitable in awaitables), return_exceptions=True)) + + +async def _insert_users( + prisma_client: PrismaClient, prepared: Sequence[_PreparedUser] +) -> tuple[tuple[_PreparedUser, ...], tuple[_RowFailure, ...]]: + """Insert every row in one statement. If that fails, retry rows one at a time so the error lands on its row.""" + if not prepared: + return (), () + table: Final = _user_table(prisma_client) + payloads: Final = tuple(_user_create_payload(user) for user in prepared) + try: + await table.create_many(data=payloads) + return tuple(prepared), () + except Exception as exc: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified + verbose_proxy_logger.warning("/user/bulk_new: create_many failed, retrying rows individually - %s", exc) + outcomes: Final = await _bounded( + BULK_NEW_USER_CONCURRENCY, tuple(table.create(data=payload) for payload in payloads) + ) + return ( + tuple(user for user, outcome in zip(prepared, outcomes, strict=True) if not isinstance(outcome, BaseException)), + tuple( + _RowFailure(user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome)) + for user, outcome in zip(prepared, outcomes, strict=True) + if isinstance(outcome, BaseException) + ), + ) + + +def _assignments_by_team(created: Sequence[_PreparedUser]) -> Mapping[str, tuple[_TeamAssignment, ...]]: + team_ids: Final = tuple(dict.fromkeys(team.team_id for user in created for team in user.pending.teams)) + return MappingProxyType( + { + team_id: tuple( + _TeamAssignment(user.pending.user_id, user.row.user_email, team.user_role, team.max_budget_in_team) + for user in created + for team in user.pending.teams + if team.team_id == team_id + ) + for team_id in team_ids + } + ) + + +class _MembershipData(TypedDict): + team_id: ReadOnly[str] + user_id: ReadOnly[str] + budget_id: ReadOnly[str | None] + + +class _RosterData(TypedDict): + members_with_roles: ReadOnly[str] + + +class _TeamsData(TypedDict): + teams: ReadOnly[tuple[str, ...]] + + +def _default_member_budget_id(team: LiteLLM_TeamTable) -> str | None: + metadata: Final = ( + _JSON_OBJECT.validate_python( + team.metadata # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_TeamTable.metadata is a bare dict; validated by the adapter + ) + if team.metadata # pyright: ignore[reportUnknownMemberType] # same bare dict + else None + ) + budget_id: Final = metadata.get("team_member_budget_id") if metadata is not None else None + return budget_id if isinstance(budget_id, str) else None + + +def _team_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamTable]": + return tx.litellm_teamtable # pyright: ignore[reportReturnType] # TableActions widens the generated inputs to Mapping, as the repositories do + + +def _membership_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamMembership]": + return tx.litellm_teammembership # pyright: ignore[reportReturnType] # TableActions widens the generated inputs to Mapping, as the repositories do + + +async def _write_team_roster( + prisma_client: PrismaClient, + team: LiteLLM_TeamTable, + members: Sequence[_TeamAssignment], + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str, +) -> _TeamWrite: + """Add every new member to one team under its advisory lock: one roster rewrite and one membership insert.""" + try: + async with prisma_client.tx() as tx: + await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team.team_id) + roster: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, team.team_id) + if roster is None: + raise ValueError(f"Team id={team.team_id} does not exist") + already_present: Final = frozenset(member.user_id for member in roster if member.user_id) + new_members: Final = tuple(member for member in members if member.user_id not in already_present) + budget_ids: Final = tuple( + [ # mutable-ok: budgets are created one at a time on the transaction's single connection + await _resolve_member_budget_id( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + max_budget_in_team=member.max_budget_in_team, + allowed_models=team.default_team_member_models or None, + budget_duration=None, + default_team_budget_id=_default_member_budget_id(team), + tx=tx, # pyright: ignore[reportArgumentType] # MemberWriteTx lags the generated Prisma signatures, same as /team/member_add + ) + for member in new_members + ] + ) + await _membership_tx_db(tx).create_many( + data=tuple( + _MembershipData(team_id=team.team_id, user_id=member.user_id, budget_id=budget_id) + for member, budget_id in zip(new_members, budget_ids, strict=True) + ), + skip_duplicates=True, + ) + after: Final = ( + *roster, + *(Member(user_id=m.user_id, user_email=m.user_email, role=m.role) for m in new_members), + ) + await _team_tx_db(tx).update( + where={"team_id": team.team_id}, # mutable-ok: Prisma query filters are dict-shaped + data=_RosterData(members_with_roles=json.dumps(tuple(member.model_dump() for member in after))), + ) + return _TeamWrite( + team_id=team.team_id, + after=after, + added=frozenset(member.user_id for member in new_members), + failed=MappingProxyType({}), + ) + except Exception as exc: # noqa: BLE001 # the team write failure is reported on each affected row + verbose_proxy_logger.exception("/user/bulk_new: failed to add members to team %s - %s", team.team_id, exc) + message: Final = f"Failed to add user to team {team.team_id}: {_error_message(exc)}" + return _TeamWrite( + team_id=team.team_id, + after=(), + added=frozenset(), + failed=MappingProxyType({member.user_id: message for member in members}), + ) + + +async def _detach_failed_teams( + prisma_client: PrismaClient, created: Sequence[_PreparedUser], writes: Mapping[str, _TeamWrite] +) -> None: + """Users are inserted with `teams` already set; drop the teams whose roster write did not take them.""" + table: Final = _user_table(prisma_client) + updates: Final = tuple( + table.update( + where={"user_id": user.row.user_id}, # mutable-ok: Prisma query filters are dict-shaped + data=_TeamsData(teams=landed), + ) + for user in created + if (landed := _row_teams(user, writes)[0]) != tuple(team.team_id for team in user.pending.teams) + ) + for outcome in await _bounded(BULK_NEW_USER_CONCURRENCY, updates): + if isinstance(outcome, BaseException): + verbose_proxy_logger.warning("/user/bulk_new: could not detach failed teams from user - %s", outcome) + + +async def _publish_team_writes(writes: Sequence[_TeamWrite], user_api_key_cache: "UserApiKeyCache") -> None: + prometheus_logger: Final = PrometheusLogger.get_instance() + for write in writes: + if prometheus_logger is None or not write.added: + continue + try: + prometheus_logger.set_team_members_metric( + LiteLLM_TeamTable( + team_id=write.team_id, + members_with_roles=write.after, # pyright: ignore[reportArgumentType] # pydantic coerces the tuple into the declared list + ) + ) + except Exception as exc: # noqa: BLE001 # metrics are best-effort and must not fail the request + verbose_proxy_logger.debug("Prometheus: failed to emit team members metric: %s", exc) + evictions: Final = await _bounded( + BULK_NEW_USER_CONCURRENCY, + tuple( + invalidate_team_member_spend_state( + user_id=user_id, team_id=write.team_id, user_api_key_cache=user_api_key_cache + ) + for write in writes + for user_id in write.added + ), + ) + for eviction in evictions: + if isinstance(eviction, BaseException): + verbose_proxy_logger.warning("/user/bulk_new: cache eviction failed - %s", eviction) + + +_KEY_FIELDS: Final = MappingProxyType( + { + name: True + for name in ( + "user_id", + "team_id", + "duration", + "key_alias", + "models", + "metadata", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + "allowed_cache_controls", + "model_max_budget", + "model_rpm_limit", + "model_tpm_limit", + "mcp_rpm_limit", + "tag_rpm_limit", + "guardrails", + "policies", + "prompts", + "object_permission_id", + ) + } +) + + +async def _generate_key(prepared: _PreparedUser, generate_key: KeyGenerator) -> str: + response: Final = _KEY_RESPONSE.validate_python( + await generate_key( + request_type="key", table_name="key", **prepared.row.model_dump(include=_KEY_FIELDS, exclude_none=True) + ) + ) + return response.token + + +async def _add_to_organizations( + prepared: _PreparedUser, organizations: Sequence[str], user_api_key_dict: UserAPIKeyAuth +) -> None: + for organization_id in organizations: + await organization_member_add( + data=OrganizationMemberAddRequest( + organization_id=organization_id, + member=OrgMember(user_id=prepared.row.user_id, role=LitellmUserRoles.INTERNAL_USER), + ), + http_request=Request(scope={"type": "http", "path": "/user/bulk_new"}), # mutable-ok: ASGI scopes are dicts + user_api_key_dict=user_api_key_dict, + ) + + +async def _run_per_user( + created: Sequence[_PreparedUser], + select: Callable[[_PreparedUser], bool], + action: Callable[[_PreparedUser], Awaitable[_T]], +) -> Mapping[str, _T | BaseException]: + chosen: Final = tuple(user for user in created if select(user)) + outcomes: Final = await _bounded(BULK_NEW_USER_CONCURRENCY, tuple(action(user) for user in chosen)) + return MappingProxyType({user.row.user_id: outcome for user, outcome in zip(chosen, outcomes, strict=True)}) + + +async def _write_audit_logs( + prisma_client: PrismaClient, + created: Sequence[_PreparedUser], + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str, +) -> None: + if not created: + return + created_ids: Final = sorted(user.row.user_id for user in created) + created_filter: Final = {"user_id": {"in": created_ids}} # mutable-ok: Prisma query filters are dict-shaped + rows: Final = await _user_table(prisma_client).find_many(where=created_filter) + outcomes: Final = await _bounded( + BULK_NEW_USER_CONCURRENCY, + tuple( + UserManagementEventHooks.create_internal_user_audit_log( + user_id=row.user_id, + action="created", + litellm_changed_by=user_api_key_dict.user_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + before_value=None, + after_value=row.model_dump_json(exclude_none=True), + ) + for row in rows + ), + ) + for outcome in outcomes: + if isinstance(outcome, BaseException): + verbose_proxy_logger.warning("Unable to create audit log for user on `/user/bulk_new` - %s", outcome) + + +def _row_teams(prepared: _PreparedUser, writes: Mapping[str, _TeamWrite]) -> tuple[tuple[str, ...], tuple[str, ...]]: + """Split a user's requested teams into the ones they landed in and the errors for the ones they did not.""" + requested: Final = tuple(team.team_id for team in prepared.pending.teams) + return ( + tuple(team_id for team_id in requested if prepared.row.user_id in writes[team_id].added), + tuple( + writes[team_id].failed[prepared.row.user_id] + for team_id in requested + if prepared.row.user_id in writes[team_id].failed + ), + ) + + +def _to_result(created: _CreatedUser) -> UserCreateResult: + return UserCreateResult( + user_id=created.prepared.row.user_id, + user_email=created.prepared.row.user_email, + success=True, + teams=created.teams, + key=created.key, + error="; ".join(created.errors) if created.errors else None, + ) + + +def _failure_result(failure: _RowFailure) -> UserCreateResult: + return UserCreateResult(user_id=failure.user_id, user_email=failure.user_email, success=False, error=failure.error) + + +async def bulk_create_users( + users: Sequence[BulkNewUserItem], + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, + license_check: LicenseCheck, + litellm_proxy_admin_name: str, + user_api_key_cache: "UserApiKeyCache", + generate_key: KeyGenerator = generate_key_helper_fn, +) -> BulkNewUserResponse: + """Create every valid row in `users`; rows that fail validation or a write are reported, not raised. + + Raises `HTTPException(403)` only when the whole batch would push the deployment over its license seat limit. + """ + pending, request_failures = _partition_rows(users, user_api_key_dict) + existing_ids, existing_emails = await _existing_user_conflicts(prisma_client, pending) + teams, team_errors = await _unusable_teams(prisma_client, pending, user_api_key_dict) + db_failures: Final = tuple( + failure + for user in pending + if (failure := _db_failure(user, existing_ids, existing_emails, team_errors)) is not None + ) + failed_indexes: Final = frozenset(failure.index for failure in db_failures) + creatable: Final = tuple(user for user in pending if user.index not in failed_indexes) + + billable_users: Final = await UserRepository(prisma_client).count_billable_users() + if creatable and license_check.is_over_limit(total_users=billable_users + len(creatable)): + raise HTTPException( + status_code=403, + detail="License is over limit. Please contact support@berri.ai to upgrade your license.", + ) + + prepared_outcomes: Final = tuple([await _prepare_user(user, prisma_client) for user in creatable]) + prepare_failures: Final = tuple(o for o in prepared_outcomes if isinstance(o, _RowFailure)) + created, insert_failures = await _insert_users( + prisma_client, tuple(o for o in prepared_outcomes if isinstance(o, _PreparedUser)) + ) + + team_writes: Final = MappingProxyType( + { + team_id: await _write_team_roster( + prisma_client, teams[team_id], members, user_api_key_dict, litellm_proxy_admin_name + ) + for team_id, members in _assignments_by_team(created).items() + } + ) + await _detach_failed_teams(prisma_client, created, team_writes) + await _publish_team_writes(tuple(team_writes.values()), user_api_key_cache) + + keys: Final = await _run_per_user( + created, lambda user: user.pending.request.auto_create_key, lambda user: _generate_key(user, generate_key) + ) + org_outcomes: Final = await _run_per_user( + created, + lambda user: bool(user.row.organizations), + lambda user: _add_to_organizations(user, user.row.organizations or (), user_api_key_dict), + ) + await _write_audit_logs(prisma_client, created, user_api_key_dict, litellm_proxy_admin_name) + + def finish(prepared: _PreparedUser) -> _CreatedUser: + landed, team_failures = _row_teams(prepared, team_writes) + key_outcome: Final = keys.get(prepared.row.user_id) + org_outcome: Final = org_outcomes.get(prepared.row.user_id) + return _CreatedUser( + prepared=prepared, + teams=landed, + key=key_outcome if isinstance(key_outcome, str) else None, + errors=( + *team_failures, + *( + (f"Failed to create key: {_error_message(key_outcome)}",) + if isinstance(key_outcome, BaseException) + else () + ), + *( + (f"Failed to add user to organizations: {_error_message(org_outcome)}",) + if isinstance(org_outcome, BaseException) + else () + ), + ), + ) + + failures: Final = MappingProxyType( + { + failure.index: _failure_result(failure) + for failure in (*request_failures, *db_failures, *prepare_failures, *insert_failures) + } + ) + successes_by_index: Final = MappingProxyType({user.pending.index: _to_result(finish(user)) for user in created}) + results: Final = tuple( + failures[index] if index in failures else successes_by_index[index] for index in range(len(users)) + ) + successes: Final = sum(1 for result in results if result.success) + return BulkNewUserResponse( + results=results, + total_requested=len(users), + successful_creations=successes, + failed_creations=len(users) - successes, + ) diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index 6973f1d1f12..a7e5c6d2916 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -1,15 +1,18 @@ from collections.abc import Mapping from typing import Any, Final, Literal -from pydantic import BaseModel, field_validator +from pydantic import BaseModel, Field, field_validator from typing_extensions import ReadOnly, TypedDict from litellm.proxy._types import ( LiteLLM_UserTableWithKeyCount, + NewUserRequest, UpdateUserRequest, UpdateUserRequestNoUserIDorEmail, ) +MAX_BULK_NEW_USERS: Final = 500 + class InsensitiveContains(TypedDict): contains: ReadOnly[str] @@ -83,3 +86,38 @@ class BulkUpdateUserResponse(BaseModel): total_requested: int successful_updates: int failed_updates: int + + +class BulkNewUserItem(NewUserRequest): + """One row of `/user/bulk_new`: the `/user/new` body, with keys opt-in and invite emails unsupported.""" + + auto_create_key: bool = False + + @field_validator("send_invite_email") + @classmethod + def reject_invite_email(cls, value: bool | None) -> bool | None: + if value: + raise ValueError("send_invite_email is not supported on /user/bulk_new; invite users separately") + return value + + +class BulkNewUserRequest(BaseModel): + users: tuple[BulkNewUserItem, ...] = Field(min_length=1, max_length=MAX_BULK_NEW_USERS) + + +class UserCreateResult(BaseModel): + """Outcome for one row of `/user/bulk_new`. `teams` lists the teams the user was actually added to.""" + + user_id: str | None = None + user_email: str | None = None + success: bool + teams: tuple[str, ...] | None = None + key: str | None = None + error: str | None = None + + +class BulkNewUserResponse(BaseModel): + results: tuple[UserCreateResult, ...] + total_requested: int + successful_creations: int + failed_creations: int diff --git a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py new file mode 100644 index 00000000000..c9f0a0e4ffe --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py @@ -0,0 +1,347 @@ +import json +from contextlib import asynccontextmanager +from typing import Final + +import pytest +from fastapi import HTTPException +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.caching.caching import DualCache +from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, UserAPIKeyAuth +from litellm.proxy.management_helpers.bulk_user_creation import bulk_create_users +from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( + BulkNewUserItem, + BulkNewUserRequest, +) + +ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) +INTERNAL: Final = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER) + + +class _UserRow(BaseModel): + model_config = ConfigDict(extra="allow") + + user_id: str + user_email: str | None = None + user_role: str | None = None + teams: list[str] = [] + max_budget: float | None = None + + +class _UserTable: + """Enough of the Prisma user table for the bulk path: set lookups, one create_many and per-row fallbacks.""" + + def __init__(self, fail_ids: frozenset[str] = frozenset()) -> None: + self.rows: dict[str, _UserRow] = {} + self.fail_ids = fail_ids + self.create_many_calls = 0 + + async def count(self, where: object = None) -> int: + return 0 if where is not None else len(self.rows) + + async def find_many(self, where: dict[str, dict[str, object]]) -> list[_UserRow]: + if "user_id" in where: + wanted = where["user_id"]["in"] + return [row for row in self.rows.values() if row.user_id in wanted] + wanted_emails = {str(e).lower() for e in where["user_email"]["in"]} + return [row for row in self.rows.values() if (row.user_email or "").lower() in wanted_emails] + + async def create(self, data: dict[str, object]) -> _UserRow: + row = _UserRow.model_validate(data) + if row.user_id in self.fail_ids or row.user_id in self.rows: + raise RuntimeError(f"insert failed for {row.user_id}") + self.rows[row.user_id] = row + return row + + async def create_many(self, data: list[dict[str, object]]) -> int: + self.create_many_calls += 1 + rows = [_UserRow.model_validate(d) for d in data] + if any(row.user_id in self.fail_ids for row in rows): + raise RuntimeError("batch insert failed") + for row in rows: + self.rows[row.user_id] = row + return len(rows) + + async def update(self, where: dict[str, str], data: dict[str, object]) -> _UserRow: + row = self.rows[where["user_id"]] + updated = _UserRow.model_validate({**row.model_dump(), **data}) + self.rows[row.user_id] = updated + return updated + + +class _TeamTable: + def __init__(self, teams: list[LiteLLM_TeamTable]) -> None: + self.rows = {team.team_id: team for team in teams} + self.update_calls = 0 + + async def find_many(self, where: dict[str, dict[str, list[str]]]) -> list[LiteLLM_TeamTable]: + return [self.rows[team_id] for team_id in where["team_id"]["in"] if team_id in self.rows] + + async def update(self, where: dict[str, str], data: dict[str, str]) -> LiteLLM_TeamTable: + self.update_calls += 1 + team = self.rows[where["team_id"]] + team.members_with_roles = [Member(**m) for m in json.loads(data["members_with_roles"])] + return team + + +class _MembershipTable: + def __init__(self) -> None: + self.rows: list[dict[str, object]] = [] + + async def create_many(self, data: list[dict[str, object]], skip_duplicates: bool = False) -> int: + self.rows.extend(data) + return len(data) + + +class _Tx: + def __init__(self, db: "_Db") -> None: + self.litellm_teamtable = db.litellm_teamtable + self.litellm_teammembership = db.litellm_teammembership + self.locks: list[str] = [] + + async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]: + if "pg_advisory_xact_lock" in sql: + self.locks.append(str(args[0])) + return [] + team = self.litellm_teamtable.rows.get(str(args[0])) + if team is None: + return [] + return [{"members_with_roles": [m.model_dump() for m in team.members_with_roles]}] + + +class _Db: + def __init__(self, teams: list[LiteLLM_TeamTable], fail_ids: frozenset[str] = frozenset()) -> None: + self.litellm_usertable = _UserTable(fail_ids) + self.litellm_teamtable = _TeamTable(teams) + self.litellm_teammembership = _MembershipTable() + + +class _FakePrisma: + def __init__(self, teams: list[LiteLLM_TeamTable] | None = None, fail_ids: frozenset[str] = frozenset()) -> None: + self.db = _Db(teams or [], fail_ids) + self.tx_count = 0 + self.locks: list[str] = [] + + def jsonify_object(self, data: dict[str, object]) -> dict[str, object]: + return data + + @asynccontextmanager + async def tx(self): + self.tx_count += 1 + tx = _Tx(self.db) + yield tx + self.locks.extend(tx.locks) + + +class _License: + def __init__(self, max_users: int | None = None) -> None: + self.max_users = max_users + self.seen: list[int] = [] + + def is_over_limit(self, total_users: int) -> bool: + self.seen.append(total_users) + return self.max_users is not None and total_users > self.max_users + + +def _team(team_id: str, members: list[Member] | None = None) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=team_id, members_with_roles=members or []) + + +async def _no_keys(**kwargs: object) -> dict[str, object]: + raise AssertionError(f"key generation was not requested: {kwargs}") + + +async def _run(prisma, users, caller=ADMIN, license=None, generate_key=_no_keys): + return await bulk_create_users( + users=[BulkNewUserItem(**u) for u in users], + user_api_key_dict=caller, + prisma_client=prisma, + license_check=license or _License(), + litellm_proxy_admin_name="default_user_id", + user_api_key_cache=DualCache(), + generate_key=generate_key, + ) + + +@pytest.mark.asyncio +async def test_creates_users_and_team_membership_in_every_store(): + prisma = _FakePrisma(teams=[_team("t1", [Member(user_id="existing", role="admin")]), _team("t2")]) + response = await _run( + prisma, + [ + {"user_id": "u1", "user_email": "a@example.com", "teams": ["t1", "t2"], "max_budget": 50}, + {"user_id": "u2", "user_email": "b@example.com", "teams": ["t1"]}, + {"user_id": "u3", "user_email": "c@example.com"}, + ], + ) + + assert (response.total_requested, response.successful_creations, response.failed_creations) == (3, 3, 0) + assert [r.user_id for r in response.results] == ["u1", "u2", "u3"] + assert all(r.success and r.key is None and r.error is None for r in response.results) + assert [r.teams for r in response.results] == [("t1", "t2"), ("t1",), ()] + + users = prisma.db.litellm_usertable.rows + assert users["u1"].teams == ["t1", "t2"] and users["u1"].max_budget == 50 + assert users["u2"].teams == ["t1"] and users["u3"].teams == [] + assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["existing", "u1", "u2"] + assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t2"].members_with_roles] == ["u1"] + assert sorted((m["team_id"], m["user_id"]) for m in prisma.db.litellm_teammembership.rows) == [ + ("t1", "u1"), + ("t1", "u2"), + ("t2", "u1"), + ] + + +@pytest.mark.asyncio +async def test_one_insert_and_one_locked_write_per_team(): + prisma = _FakePrisma(teams=[_team("t1"), _team("t2")]) + await _run( + prisma, + [{"user_id": f"u{i}", "teams": ["t1"] if i % 2 else ["t1", "t2"]} for i in range(20)], + ) + + assert prisma.db.litellm_usertable.create_many_calls == 1 + assert prisma.tx_count == 2 + assert sorted(prisma.locks) == ["t1", "t2"] + assert prisma.db.litellm_teamtable.update_calls == 2 + assert len(prisma.db.litellm_teamtable.rows["t1"].members_with_roles) == 20 + assert len(prisma.db.litellm_teamtable.rows["t2"].members_with_roles) == 10 + + +@pytest.mark.asyncio +async def test_bad_rows_fail_alone_and_good_rows_still_land(): + prisma = _FakePrisma(teams=[_team("t1")]) + prisma.db.litellm_usertable.rows["taken"] = _UserRow(user_id="taken", user_email="Taken@Example.com") + response = await _run( + prisma, + [ + {"user_id": "u1", "user_email": "a@example.com", "teams": ["t1"]}, + {"user_id": "u2", "user_email": "A@EXAMPLE.COM"}, + {"user_id": "u1", "user_email": "z@example.com"}, + {"user_id": "u3", "user_email": "taken@example.com"}, + {"user_id": "taken"}, + {"user_id": "u4", "teams": ["missing"]}, + {"user_id": "u5", "teams": ["t1", "missing"]}, + {"user_id": "u6", "budget_duration": "not-a-duration"}, + {"user_id": "u7", "user_email": "ok@example.com", "teams": ["t1"]}, + ], + ) + + assert [r.success for r in response.results] == [True, False, False, False, False, False, False, False, True] + assert (response.successful_creations, response.failed_creations) == (2, 7) + errors = [r.error for r in response.results] + assert "Duplicate user_email" in errors[1] + assert "Duplicate user_id" in errors[2] + assert "already exists" in errors[3] and "already exists" in errors[4] + assert "missing" in errors[5] and "does not exist" in errors[5] + assert "missing" in errors[6] + assert errors[7] is not None + + assert set(prisma.db.litellm_usertable.rows) == {"taken", "u1", "u7"} + assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1", "u7"] + + +@pytest.mark.asyncio +async def test_insert_failure_falls_back_to_per_row_and_reports_only_that_row(): + prisma = _FakePrisma(teams=[_team("t1")], fail_ids=frozenset({"u2"})) + response = await _run( + prisma, + [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2", "teams": ["t1"]}, {"user_id": "u3"}], + ) + + assert [r.success for r in response.results] == [True, False, True] + assert "insert failed for u2" in (response.results[1].error or "") + assert set(prisma.db.litellm_usertable.rows) == {"u1", "u3"} + assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1"] + + +@pytest.mark.asyncio +async def test_team_write_failure_keeps_user_and_reports_it_on_the_row(): + prisma = _FakePrisma(teams=[_team("t1"), _team("t2")]) + + async def explode(where, data): + raise RuntimeError("roster write failed") + + prisma.db.litellm_teamtable.update = explode + response = await _run(prisma, [{"user_id": "u1", "teams": ["t1", "t2"]}]) + + result = response.results[0] + assert result.success is True + assert result.teams == () + assert "t1" in (result.error or "") and "roster write failed" in (result.error or "") + assert prisma.db.litellm_usertable.rows["u1"].teams == [] + assert (response.successful_creations, response.failed_creations) == (1, 0) + + +@pytest.mark.asyncio +async def test_keys_are_opt_in_per_row(): + prisma = _FakePrisma() + calls: list[dict[str, object]] = [] + + async def generate_key(**kwargs: object) -> dict[str, object]: + calls.append(kwargs) + return {"token": f"sk-{kwargs['user_id']}"} + + response = await _run( + prisma, + [ + {"user_id": "u1"}, + {"user_id": "u2", "auto_create_key": True, "models": ["gpt-4o"], "key_alias": "u2-key"}, + {"user_id": "u3", "auto_create_key": False}, + ], + generate_key=generate_key, + ) + + assert [r.key for r in response.results] == [None, "sk-u2", None] + assert len(calls) == 1 + assert calls[0]["user_id"] == "u2" and calls[0]["table_name"] == "key" + assert calls[0]["models"] == ("gpt-4o",) and calls[0]["key_alias"] == "u2-key" + assert set(prisma.db.litellm_usertable.rows) == {"u1", "u2", "u3"} + + +@pytest.mark.asyncio +async def test_non_admin_cannot_create_admin_users_but_other_rows_proceed(): + prisma = _FakePrisma() + response = await _run( + prisma, + [{"user_id": "u1", "user_role": "proxy_admin"}, {"user_id": "u2", "user_role": "internal_user"}], + caller=INTERNAL, + ) + + assert [r.success for r in response.results] == [False, True] + assert "Only proxy admins" in (response.results[0].error or "") + assert set(prisma.db.litellm_usertable.rows) == {"u2"} + + +@pytest.mark.asyncio +async def test_license_is_checked_once_against_the_whole_batch(): + prisma = _FakePrisma() + prisma.db.litellm_usertable.rows["existing"] = _UserRow(user_id="existing") + license = _License(max_users=3) + + with pytest.raises(HTTPException) as exc: + await _run(prisma, [{"user_id": f"u{i}"} for i in range(3)], license=license) + + assert exc.value.status_code == 403 + assert license.seen == [4] + assert set(prisma.db.litellm_usertable.rows) == {"existing"} + + ok = await _run(prisma, [{"user_id": f"u{i}"} for i in range(2)], license=license) + assert ok.successful_creations == 2 + + resend = await _run(prisma, [{"user_id": f"u{i}"} for i in range(2)], license=license) + assert [r.success for r in resend.results] == [False, False] + assert all("already exists" in (r.error or "") for r in resend.results) + assert license.seen == [4, 3] + assert set(prisma.db.litellm_usertable.rows) == {"existing", "u0", "u1"} + + +def test_request_rejects_empty_oversized_and_invite_rows(): + with pytest.raises(ValidationError): + BulkNewUserRequest(users=[]) + with pytest.raises(ValidationError): + BulkNewUserRequest(users=[{"user_email": f"{i}@example.com"} for i in range(501)]) + with pytest.raises(ValidationError, match="send_invite_email"): + BulkNewUserItem(user_email="a@example.com", send_invite_email=True) + assert len(BulkNewUserRequest(users=[{"user_email": f"{i}@example.com"} for i in range(500)]).users) == 500 + assert BulkNewUserItem(user_email="a@example.com").auto_create_key is False From 24a1d772b6d2b0d8396145c12bd3ac638687a460 Mon Sep 17 00:00:00 2001 From: ryan Date: Mon, 14 Sep 2026 01:48:44 +0000 Subject: [PATCH 044/107] fix(proxy): keep key policy fields and reconcile committed rows in /user/bulk_new Rows opting into auto_create_key lost blocked, permissions, aliases, config, agent_id, budget_fallbacks and budget_limits before reaching the key helper. When create_many commits but the response is lost, re-read which ids landed and retry only the rest so committed rows report success and get their teams. Regenerate schema.d.ts and allowlist the endpoint in the Terraform audit. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_helpers/bulk_user_creation.py | 39 ++- .../endpointaudit/coverage_allowlist.txt | 1 + .../test_bulk_user_creation.py | 48 +++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 242 +++++++++++++++++- 4 files changed, 315 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index 301629e6c2a..7558a862993 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -124,6 +124,13 @@ class _UserRow(BaseModel): prompts: tuple[str, ...] | None = None duration: str | None = None key_alias: str | None = None + aliases: Mapping[str, object] | None = None + config: Mapping[str, object] | None = None + permissions: Mapping[str, object] | None = None + blocked: bool | None = None + agent_id: str | None = None + budget_fallbacks: Mapping[str, tuple[str, ...]] | None = None + budget_limits: tuple[Mapping[str, object], ...] | None = None organizations: tuple[str, ...] | None = None @@ -428,16 +435,26 @@ async def _insert_users( return tuple(prepared), () except Exception as exc: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified verbose_proxy_logger.warning("/user/bulk_new: create_many failed, retrying rows individually - %s", exc) + landed_rows: Final = await table.find_many( + where={"user_id": {"in": [payload["user_id"] for payload in payloads]}} # mutable-ok: Prisma filter + ) + landed: Final = frozenset(row.user_id for row in landed_rows) + retried: Final = tuple(user for user in prepared if user.row.user_id not in landed) outcomes: Final = await _bounded( - BULK_NEW_USER_CONCURRENCY, tuple(table.create(data=payload) for payload in payloads) + BULK_NEW_USER_CONCURRENCY, tuple(table.create(data=_user_create_payload(user)) for user in retried) + ) + failed: Final = MappingProxyType( + { + user.row.user_id: _RowFailure( + user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome) + ) + for user, outcome in zip(retried, outcomes, strict=True) + if isinstance(outcome, BaseException) + } ) return ( - tuple(user for user, outcome in zip(prepared, outcomes, strict=True) if not isinstance(outcome, BaseException)), - tuple( - _RowFailure(user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome)) - for user, outcome in zip(prepared, outcomes, strict=True) - if isinstance(outcome, BaseException) - ), + tuple(user for user in prepared if user.row.user_id not in failed), + tuple(failed.values()), ) @@ -606,9 +623,17 @@ _KEY_FIELDS: Final = MappingProxyType( for name in ( "user_id", "team_id", + "agent_id", "duration", "key_alias", "models", + "aliases", + "config", + "permissions", + "blocked", + "spend", + "budget_fallbacks", + "budget_limits", "metadata", "max_parallel_requests", "tpm_limit", diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 6bc8947e89f..6f85df98850 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -84,6 +84,7 @@ POST /team/{team_id}/member/{user_id}/reset_spend POST /team/key/bulk_update POST /team/permissions_bulk_update POST /team/{team_id}/disable_logging +POST /user/bulk_new POST /user/bulk_update # Alternate method or path for functionality the provider already manages elsewhere diff --git a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py index c9f0a0e4ffe..7ff9b037cc1 100644 --- a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py @@ -31,9 +31,10 @@ class _UserRow(BaseModel): class _UserTable: """Enough of the Prisma user table for the bulk path: set lookups, one create_many and per-row fallbacks.""" - def __init__(self, fail_ids: frozenset[str] = frozenset()) -> None: + def __init__(self, fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False) -> None: self.rows: dict[str, _UserRow] = {} self.fail_ids = fail_ids + self.commit_then_drop = commit_then_drop self.create_many_calls = 0 async def count(self, where: object = None) -> int: @@ -60,6 +61,8 @@ class _UserTable: raise RuntimeError("batch insert failed") for row in rows: self.rows[row.user_id] = row + if self.commit_then_drop: + raise ConnectionError("connection reset after commit") return len(rows) async def update(self, where: dict[str, str], data: dict[str, object]) -> _UserRow: @@ -110,15 +113,22 @@ class _Tx: class _Db: - def __init__(self, teams: list[LiteLLM_TeamTable], fail_ids: frozenset[str] = frozenset()) -> None: - self.litellm_usertable = _UserTable(fail_ids) + def __init__( + self, teams: list[LiteLLM_TeamTable], fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False + ) -> None: + self.litellm_usertable = _UserTable(fail_ids, commit_then_drop) self.litellm_teamtable = _TeamTable(teams) self.litellm_teammembership = _MembershipTable() class _FakePrisma: - def __init__(self, teams: list[LiteLLM_TeamTable] | None = None, fail_ids: frozenset[str] = frozenset()) -> None: - self.db = _Db(teams or [], fail_ids) + def __init__( + self, + teams: list[LiteLLM_TeamTable] | None = None, + fail_ids: frozenset[str] = frozenset(), + commit_then_drop: bool = False, + ) -> None: + self.db = _Db(teams or [], fail_ids, commit_then_drop) self.tx_count = 0 self.locks: list[str] = [] @@ -255,6 +265,17 @@ async def test_insert_failure_falls_back_to_per_row_and_reports_only_that_row(): assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1"] +@pytest.mark.asyncio +async def test_insert_that_committed_but_lost_its_response_still_counts_as_created(): + prisma = _FakePrisma(teams=[_team("t1")], commit_then_drop=True) + response = await _run(prisma, [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2"}]) + + assert [r.success for r in response.results] == [True, True] + assert [r.error for r in response.results] == [None, None] + assert set(prisma.db.litellm_usertable.rows) == {"u1", "u2"} + assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1"] + + @pytest.mark.asyncio async def test_team_write_failure_keeps_user_and_reports_it_on_the_row(): prisma = _FakePrisma(teams=[_team("t1"), _team("t2")]) @@ -286,7 +307,17 @@ async def test_keys_are_opt_in_per_row(): prisma, [ {"user_id": "u1"}, - {"user_id": "u2", "auto_create_key": True, "models": ["gpt-4o"], "key_alias": "u2-key"}, + { + "user_id": "u2", + "auto_create_key": True, + "models": ["gpt-4o"], + "key_alias": "u2-key", + "blocked": True, + "permissions": {"get_spend_routes": True}, + "aliases": {"fast": "gpt-4o"}, + "config": {"tier": "gold"}, + "budget_fallbacks": {"gpt-4o": ["gpt-4o-mini"]}, + }, {"user_id": "u3", "auto_create_key": False}, ], generate_key=generate_key, @@ -296,6 +327,11 @@ async def test_keys_are_opt_in_per_row(): assert len(calls) == 1 assert calls[0]["user_id"] == "u2" and calls[0]["table_name"] == "key" assert calls[0]["models"] == ("gpt-4o",) and calls[0]["key_alias"] == "u2-key" + assert calls[0]["blocked"] is True + assert calls[0]["permissions"] == {"get_spend_routes": True} + assert calls[0]["aliases"] == {"fast": "gpt-4o"} + assert calls[0]["config"] == {"tier": "gold"} + assert calls[0]["budget_fallbacks"] == {"gpt-4o": ("gpt-4o-mini",)} assert set(prisma.db.litellm_usertable.rows) == {"u1", "u2", "u3"} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..f63bed3c3ec 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16478,6 +16478,53 @@ export interface paths { patch?: never; trace?: never; }; + "/user/bulk_new": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Bulk New User + * @description Create up to 500 internal users in one request, optionally adding each one to teams. + * + * Every entry in `users` takes the same fields as `/user/new`, with two differences: `auto_create_key` + * defaults to `false` (opt in per user to also get a virtual key back) and `send_invite_email` is not + * supported. Rows are validated together (duplicate ids or emails, unknown teams, roles the caller may not + * grant), inserted in one statement, and each referenced team is written once for all of its new members. + * + * Rows fail independently: a bad row is reported in `results` with `success: false` and an `error`, and the + * other rows still get created. A user that was created but could not be added to one of its teams is + * reported with `success: true`, `teams` listing where they did land, and `error` naming the failed team. + * The whole request is rejected with 403 only if creating the valid rows would exceed the license seat limit. + * + * Usage Example + * + * ```shell + * curl -X POST "http://localhost:4000/user/bulk_new" \ + * -H "Content-Type: application/json" \ + * -H "Authorization: Bearer sk-1234" \ + * -d '{ + * "users": [ + * {"user_email": "a@example.com", "user_role": "internal_user", "teams": ["team-1"]}, + * {"user_email": "b@example.com", "user_role": "internal_user", "auto_create_key": true} + * ] + * }' + * ``` + * + * Returns `results` (one entry per input row, in order, with `user_id`, `user_email`, `success`, `teams`, + * `key`, `error`), `total_requested`, `successful_creations` and `failed_creations`. + */ + post: operations["bulk_new_user_user_bulk_new_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/user/bulk_update": { parameters: { query?: never; @@ -16781,7 +16828,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16933,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -24505,6 +24550,148 @@ export interface components { /** Budgets */ budgets: string[]; }; + /** + * BulkNewUserItem + * @description One row of `/user/bulk_new`: the `/user/new` body, with keys opt-in and invite emails unsupported. + */ + BulkNewUserItem: { + /** Agent Id */ + agent_id?: string | null; + /** + * Aliases + * @default {} + */ + aliases: { + [key: string]: unknown; + } | null; + /** + * Allowed Cache Controls + * @default [] + */ + allowed_cache_controls: unknown[] | null; + /** + * Auto Create Key + * @default false + */ + auto_create_key: boolean; + /** Blocked */ + blocked?: boolean | null; + /** Budget Duration */ + budget_duration?: string | null; + /** Budget Fallbacks */ + budget_fallbacks?: { + [key: string]: string[]; + } | null; + /** Budget Limits */ + budget_limits?: components["schemas"]["BudgetLimitEntry"][] | null; + /** + * Config + * @default {} + */ + config: { + [key: string]: unknown; + } | null; + /** Duration */ + duration?: string | null; + /** Guardrails */ + guardrails?: string[] | null; + /** Key Alias */ + key_alias?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** Max Parallel Requests */ + max_parallel_requests?: number | null; + /** Mcp Rpm Limit */ + mcp_rpm_limit?: { + [key: string]: number; + } | null; + /** + * Metadata + * @default {} + */ + metadata: { + [key: string]: unknown; + } | null; + /** + * Model Max Budget + * @default {} + */ + model_max_budget: { + [key: string]: unknown; + } | null; + /** Model Rpm Limit */ + model_rpm_limit?: { + [key: string]: unknown; + } | null; + /** Model Tpm Limit */ + model_tpm_limit?: { + [key: string]: unknown; + } | null; + /** + * Models + * @default [] + */ + models: unknown[] | null; + object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null; + /** Organizations */ + organizations?: string[] | null; + /** + * Permissions + * @default {} + */ + permissions: { + [key: string]: unknown; + } | null; + /** Policies */ + policies?: string[] | null; + /** Prompts */ + prompts?: string[] | null; + /** Rpm Limit */ + rpm_limit?: number | null; + /** Send Invite Email */ + send_invite_email?: boolean | null; + /** + * Spend + * @default 0 + */ + spend: number | null; + /** Sso User Id */ + sso_user_id?: string | null; + /** Tag Rpm Limit */ + tag_rpm_limit?: { + [key: string]: number; + } | null; + /** Team Id */ + team_id?: string | null; + /** Teams */ + teams?: string[] | components["schemas"]["NewUserRequestTeam"][] | null; + /** Tpm Limit */ + tpm_limit?: number | null; + /** User Alias */ + user_alias?: string | null; + /** User Email */ + user_email?: string | null; + /** User Id */ + user_id?: string | null; + /** User Role */ + user_role?: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; + }; + /** BulkNewUserRequest */ + BulkNewUserRequest: { + /** Users */ + users: components["schemas"]["BulkNewUserItem"][]; + }; + /** BulkNewUserResponse */ + BulkNewUserResponse: { + /** Failed Creations */ + failed_creations: number; + /** Results */ + results: components["schemas"]["UserCreateResult"][]; + /** Successful Creations */ + successful_creations: number; + /** Total Requested */ + total_requested: number; + }; /** * BulkTeamMemberAddRequest * @description Request for bulk team member addition @@ -39345,6 +39532,24 @@ export interface components { */ severity: "info" | "warning" | "error"; }; + /** + * UserCreateResult + * @description Outcome for one row of `/user/bulk_new`. `teams` lists the teams the user was actually added to. + */ + UserCreateResult: { + /** Error */ + error?: string | null; + /** Key */ + key?: string | null; + /** Success */ + success: boolean; + /** Teams */ + teams?: string[] | null; + /** User Email */ + user_email?: string | null; + /** User Id */ + user_id?: string | null; + }; /** * UserHeaderMapping * @description Map an incoming HTTP header to a LiteLLM user role. @@ -60631,6 +60836,39 @@ export interface operations { }; }; }; + bulk_new_user_user_bulk_new_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["BulkNewUserRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["BulkNewUserResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; bulk_user_update_user_bulk_update_post: { parameters: { query?: never; From f5e7294b46efb2b33b7184043f22f0ebf5160e2f Mon Sep 17 00:00:00 2001 From: ryan Date: Mon, 14 Sep 2026 02:12:54 +0000 Subject: [PATCH 045/107] fix(proxy): keep request identifiers out of /user/bulk_new log lines Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_helpers/bulk_user_creation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index 7558a862993..a55445634dc 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -348,7 +348,7 @@ async def _prepare_user(user: _PendingUser, prisma_client: PrismaClient) -> _Pre ) return _PreparedUser(user, _USER_ROW.validate_python(with_permission)) except Exception as exc: # noqa: BLE001 # any preparation failure is reported on this row only - verbose_proxy_logger.warning("/user/bulk_new: could not prepare user %s - %s", user.user_id, exc) + verbose_proxy_logger.warning("/user/bulk_new: could not prepare row %d - %s", user.index, type(exc).__name__) return _RowFailure(user.index, user.user_id, user.request.user_email, _error_message(exc)) @@ -560,7 +560,7 @@ async def _write_team_roster( failed=MappingProxyType({}), ) except Exception as exc: # noqa: BLE001 # the team write failure is reported on each affected row - verbose_proxy_logger.exception("/user/bulk_new: failed to add members to team %s - %s", team.team_id, exc) + verbose_proxy_logger.exception("/user/bulk_new: failed to add %d members to a team", len(members)) message: Final = f"Failed to add user to team {team.team_id}: {_error_message(exc)}" return _TeamWrite( team_id=team.team_id, From 7ec9e2a7e882114831582e9a9a1fe1fd3089f4ae Mon Sep 17 00:00:00 2001 From: ryan Date: Mon, 14 Sep 2026 02:35:36 +0000 Subject: [PATCH 046/107] fix(proxy): log /user/bulk_new failures with exc_info instead of request-derived values Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../internal_user_endpoints.py | 2 +- .../management_helpers/bulk_user_creation.py | 18 +++++++++++------- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 2624d7e373a..5f570485ed3 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -701,7 +701,7 @@ async def bulk_new_user( user_api_key_cache=user_api_key_cache, ) except Exception as e: # noqa: BLE001 # normalize every failure to the proxy exception contract - verbose_proxy_logger.exception("/user/bulk_new: Exception occured - %s", e) + verbose_proxy_logger.exception("/user/bulk_new: Exception occured") raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index a55445634dc..3b587e932d4 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -433,8 +433,8 @@ async def _insert_users( try: await table.create_many(data=payloads) return tuple(prepared), () - except Exception as exc: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified - verbose_proxy_logger.warning("/user/bulk_new: create_many failed, retrying rows individually - %s", exc) + except Exception: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified + verbose_proxy_logger.warning("/user/bulk_new: create_many failed, retrying rows individually", exc_info=True) landed_rows: Final = await table.find_many( where={"user_id": {"in": [payload["user_id"] for payload in payloads]}} # mutable-ok: Prisma filter ) @@ -585,7 +585,9 @@ async def _detach_failed_teams( ) for outcome in await _bounded(BULK_NEW_USER_CONCURRENCY, updates): if isinstance(outcome, BaseException): - verbose_proxy_logger.warning("/user/bulk_new: could not detach failed teams from user - %s", outcome) + verbose_proxy_logger.warning( + "/user/bulk_new: could not detach failed teams from user - %s", type(outcome).__name__ + ) async def _publish_team_writes(writes: Sequence[_TeamWrite], user_api_key_cache: "UserApiKeyCache") -> None: @@ -600,8 +602,8 @@ async def _publish_team_writes(writes: Sequence[_TeamWrite], user_api_key_cache: members_with_roles=write.after, # pyright: ignore[reportArgumentType] # pydantic coerces the tuple into the declared list ) ) - except Exception as exc: # noqa: BLE001 # metrics are best-effort and must not fail the request - verbose_proxy_logger.debug("Prometheus: failed to emit team members metric: %s", exc) + except Exception: # noqa: BLE001 # metrics are best-effort and must not fail the request + verbose_proxy_logger.debug("Prometheus: failed to emit team members metric", exc_info=True) evictions: Final = await _bounded( BULK_NEW_USER_CONCURRENCY, tuple( @@ -614,7 +616,7 @@ async def _publish_team_writes(writes: Sequence[_TeamWrite], user_api_key_cache: ) for eviction in evictions: if isinstance(eviction, BaseException): - verbose_proxy_logger.warning("/user/bulk_new: cache eviction failed - %s", eviction) + verbose_proxy_logger.warning("/user/bulk_new: cache eviction failed - %s", type(eviction).__name__) _KEY_FIELDS: Final = MappingProxyType( @@ -714,7 +716,9 @@ async def _write_audit_logs( ) for outcome in outcomes: if isinstance(outcome, BaseException): - verbose_proxy_logger.warning("Unable to create audit log for user on `/user/bulk_new` - %s", outcome) + verbose_proxy_logger.warning( + "Unable to create audit log for user on `/user/bulk_new` - %s", type(outcome).__name__ + ) def _row_teams(prepared: _PreparedUser, writes: Mapping[str, _TeamWrite]) -> tuple[tuple[str, ...], tuple[str, ...]]: From 4285f1dfb09f9a5ccdda626fa0cd70074fca1354 Mon Sep 17 00:00:00 2001 From: ryan Date: Mon, 14 Sep 2026 02:39:47 +0000 Subject: [PATCH 047/107] fix(proxy): do not claim rows a concurrent request inserted when /user/bulk_new create_many fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_helpers/bulk_user_creation.py | 34 +++++++++++----- .../test_bulk_user_creation.py | 39 ++++++++++++++++--- 2 files changed, 59 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index 3b587e932d4..8bd7e83d122 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -31,6 +31,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import invalidate_team_member_spend_state from litellm.proxy.auth.litellm_license import LicenseCheck from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks from litellm.proxy.management_endpoints.common_utils import ( _is_user_org_admin_for_team, # pyright: ignore[reportPrivateUsage] # same team-admin check /user/new uses @@ -433,23 +434,38 @@ async def _insert_users( try: await table.create_many(data=payloads) return tuple(prepared), () - except Exception: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified + except Exception as exc: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified verbose_proxy_logger.warning("/user/bulk_new: create_many failed, retrying rows individually", exc_info=True) - landed_rows: Final = await table.find_many( - where={"user_id": {"in": [payload["user_id"] for payload in payloads]}} # mutable-ok: Prisma filter - ) + outcome_unknown: Final = PrismaDBExceptionHandler.is_database_infrastructure_error(exc) + requested: Final = frozenset(payload["user_id"] for payload in payloads) + landed_rows: Final = await table.find_many(where={"user_id": {"in": list(requested)}}) # mutable-ok: Prisma filter landed: Final = frozenset(row.user_id for row in landed_rows) + # create_many is one INSERT: after a lost response the full set is ours, any partial set belongs to another request + if outcome_unknown and landed == requested: + return tuple(prepared), () + taken: Final = tuple(user for user in prepared if user.row.user_id in landed) retried: Final = tuple(user for user in prepared if user.row.user_id not in landed) outcomes: Final = await _bounded( BULK_NEW_USER_CONCURRENCY, tuple(table.create(data=_user_create_payload(user)) for user in retried) ) failed: Final = MappingProxyType( { - user.row.user_id: _RowFailure( - user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome) - ) - for user, outcome in zip(retried, outcomes, strict=True) - if isinstance(outcome, BaseException) + **{ + user.row.user_id: _RowFailure( + user.pending.index, + user.pending.user_id, + user.row.user_email, + f"User id={user.row.user_id} already exists", + ) + for user in taken + }, + **{ + user.row.user_id: _RowFailure( + user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome) + ) + for user, outcome in zip(retried, outcomes, strict=True) + if isinstance(outcome, BaseException) + }, } ) return ( diff --git a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py index 7ff9b037cc1..00ddd63c057 100644 --- a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py @@ -2,8 +2,10 @@ import json from contextlib import asynccontextmanager from typing import Final +import httpx import pytest from fastapi import HTTPException +from prisma.errors import UniqueViolationError from pydantic import BaseModel, ConfigDict, ValidationError from litellm.caching.caching import DualCache @@ -31,10 +33,16 @@ class _UserRow(BaseModel): class _UserTable: """Enough of the Prisma user table for the bulk path: set lookups, one create_many and per-row fallbacks.""" - def __init__(self, fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False) -> None: + def __init__( + self, + fail_ids: frozenset[str] = frozenset(), + commit_then_drop: bool = False, + raced_ids: frozenset[str] = frozenset(), + ) -> None: self.rows: dict[str, _UserRow] = {} self.fail_ids = fail_ids self.commit_then_drop = commit_then_drop + self.raced_ids = raced_ids self.create_many_calls = 0 async def count(self, where: object = None) -> int: @@ -59,10 +67,15 @@ class _UserTable: rows = [_UserRow.model_validate(d) for d in data] if any(row.user_id in self.fail_ids for row in rows): raise RuntimeError("batch insert failed") + raced = [row.user_id for row in rows if row.user_id in self.raced_ids] + if raced: + for user_id in raced: + self.rows[user_id] = _UserRow(user_id=user_id, user_email=f"{user_id}@other-request.example") + raise UniqueViolationError({}, message="Unique constraint failed on the fields: (`user_id`)") for row in rows: self.rows[row.user_id] = row if self.commit_then_drop: - raise ConnectionError("connection reset after commit") + raise httpx.ReadError("connection reset after commit") return len(rows) async def update(self, where: dict[str, str], data: dict[str, object]) -> _UserRow: @@ -114,9 +127,13 @@ class _Tx: class _Db: def __init__( - self, teams: list[LiteLLM_TeamTable], fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False + self, + teams: list[LiteLLM_TeamTable], + fail_ids: frozenset[str] = frozenset(), + commit_then_drop: bool = False, + raced_ids: frozenset[str] = frozenset(), ) -> None: - self.litellm_usertable = _UserTable(fail_ids, commit_then_drop) + self.litellm_usertable = _UserTable(fail_ids, commit_then_drop, raced_ids) self.litellm_teamtable = _TeamTable(teams) self.litellm_teammembership = _MembershipTable() @@ -127,8 +144,9 @@ class _FakePrisma: teams: list[LiteLLM_TeamTable] | None = None, fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False, + raced_ids: frozenset[str] = frozenset(), ) -> None: - self.db = _Db(teams or [], fail_ids, commit_then_drop) + self.db = _Db(teams or [], fail_ids, commit_then_drop, raced_ids) self.tx_count = 0 self.locks: list[str] = [] @@ -276,6 +294,17 @@ async def test_insert_that_committed_but_lost_its_response_still_counts_as_creat assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1"] +@pytest.mark.asyncio +async def test_user_id_taken_by_a_concurrent_request_is_not_claimed_by_this_batch(): + prisma = _FakePrisma(teams=[_team("t1")], raced_ids=frozenset({"u1"})) + response = await _run(prisma, [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2", "teams": ["t1"]}]) + + assert [r.success for r in response.results] == [False, True] + assert "User id=u1 already exists" in (response.results[0].error or "") + assert prisma.db.litellm_usertable.rows["u1"].user_email == "u1@other-request.example" + assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u2"] + + @pytest.mark.asyncio async def test_team_write_failure_keeps_user_and_reports_it_on_the_row(): prisma = _FakePrisma(teams=[_team("t1"), _team("t2")]) From 0c9e0c407e290f9b2d14b041ea8290e2f372e5a0 Mon Sep 17 00:00:00 2001 From: ryan Date: Mon, 14 Sep 2026 03:37:20 +0000 Subject: [PATCH 048/107] fix(proxy): keep a team on a /user/bulk_new row when the roster already lists that user id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_helpers/bulk_user_creation.py | 2 +- .../management_helpers/test_bulk_user_creation.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index 8bd7e83d122..2eed5121755 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -572,7 +572,7 @@ async def _write_team_roster( return _TeamWrite( team_id=team.team_id, after=after, - added=frozenset(member.user_id for member in new_members), + added=frozenset(member.user_id for member in members), failed=MappingProxyType({}), ) except Exception as exc: # noqa: BLE001 # the team write failure is reported on each affected row diff --git a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py index 00ddd63c057..16e56c1b723 100644 --- a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py @@ -220,6 +220,18 @@ async def test_creates_users_and_team_membership_in_every_store(): ] +@pytest.mark.asyncio +async def test_user_id_already_on_the_roster_keeps_the_team_and_is_not_added_twice(): + prisma = _FakePrisma(teams=[_team("t1", [Member(user_id="u1", role="user")])]) + response = await _run(prisma, [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2", "teams": ["t1"]}]) + + assert [r.success for r in response.results] == [True, True] + assert [r.teams for r in response.results] == [("t1",), ("t1",)] + assert [r.error for r in response.results] == [None, None] + assert prisma.db.litellm_usertable.rows["u1"].teams == ["t1"] + assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1", "u2"] + + @pytest.mark.asyncio async def test_one_insert_and_one_locked_write_per_team(): prisma = _FakePrisma(teams=[_team("t1"), _team("t2")]) From e41faf54a82c3443d5dbad30ad9d519d57093895 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 18:39:10 +0000 Subject: [PATCH 049/107] fix(ui): anchor guardrail lifecycle on timed entries and show not_run skip reason Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../GuardrailViewer/GuardrailViewer.test.tsx | 48 +++++++++++++++++-- .../GuardrailViewer/GuardrailViewer.tsx | 20 ++++++-- .../GuardrailViewer/__tests__/fixtures.ts | 8 ++-- 3 files changed, 62 insertions(+), 14 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx index 8cc0d186631..629383d5807 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx @@ -49,17 +49,55 @@ describe("GuardrailViewer", () => { expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); }); - it("renders not_run as NOT RUN (muted) and keeps it out of the evaluated and passed counts", () => { - const data = makeGuardrailInformation({ guardrail_status: "not_run", guardrail_mode: "pre_call" }); + it("renders not_run as NOT RUN (muted) and keeps it out of the evaluated and passed counts", async () => { + const user = userEvent.setup(); + const data = makeGuardrailInformation({ + guardrail_status: "not_run", + guardrail_mode: "pre_call", + guardrail_response: "no scannable content after message scoping", + start_time: null, + end_time: null, + duration: null, + }); renderWithProviders(); expect(screen.getByText(/0 guardrails evaluated/)).toBeInTheDocument(); expect(screen.getByText(/0 Passed/)).toHaveClass("text-muted-foreground"); expect(screen.getByText(/1 Not run/)).toBeInTheDocument(); - const badges = screen.getAllByText("NOT RUN"); - expect(badges).toHaveLength(2); - expect(badges[0]).toHaveClass("text-muted-foreground"); + const badge = screen.getByText("NOT RUN"); + expect(badge).toHaveClass("text-muted-foreground"); expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); + expect(screen.queryByText(/^T\+/)).not.toBeInTheDocument(); + + await user.click(screen.getByText("pii-rail")); + expect(screen.getByText("no scannable content after message scoping")).toBeInTheDocument(); + }); + + it("anchors the lifecycle timeline on timed entries when an untimed not_run entry sorts first", () => { + const skipped = makeGuardrailInformation({ + guardrail_name: "skipped-rail", + guardrail_status: "not_run", + guardrail_mode: "pre_call", + start_time: null, + end_time: null, + duration: null, + }); + const ran = makeGuardrailInformation({ + guardrail_name: "ran-rail", + guardrail_status: "success", + guardrail_mode: "post_call", + start_time: 1_700_000_000, + end_time: 1_700_000_000.25, + duration: 0.25, + }); + renderWithProviders(); + + expect(screen.getByText(/1 guardrail evaluated/)).toBeInTheDocument(); + expect(screen.getByText("Request received").parentElement).toHaveTextContent("T+0ms"); + expect(screen.getByText(/Post-call guardrail: ran-rail/).parentElement).toHaveTextContent("T+250ms"); + expect(screen.getByText("Response returned").parentElement).toHaveTextContent("T+251ms"); + expect(screen.queryByText(/Pre-call guardrail: skipped-rail/)).not.toBeInTheDocument(); + expect(screen.getByText("—")).toBeInTheDocument(); }); it("calculates and displays masked entity totals", async () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx index 7fb4124262b..1de0e3878b2 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx @@ -38,9 +38,9 @@ interface MatchDetail { } interface GuardrailInformation { - duration: number; - end_time: number; - start_time: number; + duration: number | null; + end_time: number | null; + start_time: number | null; guardrail_mode: string | string[] | Record | null; guardrail_name: string; guardrail_status: string; @@ -121,7 +121,8 @@ const formatMode = (mode: GuardrailInformation["guardrail_mode"]): string => { return s.replace(/_/g, "-").toUpperCase(); }; -const formatDurationMs = (seconds: number): string => { +const formatDurationMs = (seconds: number | null): string => { + if (seconds == null) return "—"; const ms = Math.round(seconds * 1000); return `${ms}ms`; }; @@ -364,8 +365,13 @@ interface TimelineEntry { outcome?: EntryOutcome; } +type TimedGuardrailInformation = GuardrailInformation & { start_time: number; end_time: number }; + +const isTimed = (e: GuardrailInformation): e is TimedGuardrailInformation => + typeof e.start_time === "number" && typeof e.end_time === "number"; + const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { - const sorted = useMemo(() => [...entries].sort((a, b) => (a.start_time ?? 0) - (b.start_time ?? 0)), [entries]); + const sorted = useMemo(() => entries.filter(isTimed).sort((a, b) => a.start_time - b.start_time), [entries]); const timeline = useMemo(() => { if (sorted.length === 0) return []; @@ -669,6 +675,10 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { )} + {outcome === "not_run" && typeof guardrailResponse === "string" && ( +

{guardrailResponse}

+ )} + {/* Provider-specific details */} {guardrailProvider === "presidio" && presidioEntities.length > 0 && (
diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/__tests__/fixtures.ts b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/__tests__/fixtures.ts index fe27428283d..ab121adf6b4 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/__tests__/fixtures.ts +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/__tests__/fixtures.ts @@ -20,13 +20,13 @@ export interface GuardrailEntity { } export interface GuardrailInformation { - duration: number; - end_time: number; - start_time: number; + duration: number | null; + end_time: number | null; + start_time: number | null; guardrail_mode: string | string[] | Record | null; guardrail_name: string; guardrail_status: string; - guardrail_response: GuardrailEntity[] | BedrockGuardrailResponse; + guardrail_response: GuardrailEntity[] | BedrockGuardrailResponse | string; masked_entity_count: Record; guardrail_usage?: Record; guardrail_cost?: number; From b1a006ea66ffb209c237fe8ab4a4df70d9806877 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 18:54:24 +0000 Subject: [PATCH 050/107] test(ui): hoist not_run guardrail fixtures out of inline call args Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../GuardrailViewer/GuardrailViewer.test.tsx | 46 +++++++++---------- 1 file changed, 22 insertions(+), 24 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx index 629383d5807..7f343211596 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx @@ -3,6 +3,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import userEvent from "@testing-library/user-event"; import { renderWithProviders, screen, waitFor } from "../../../../tests/test-utils"; import { + GuardrailInformation, makeBedrockResponse, makeEntity, makeGuardrailInformation, @@ -14,6 +15,24 @@ import GuardrailViewer from "@/components/view_logs/GuardrailViewer/GuardrailVie const PresidioPath = "@/components/view_logs/GuardrailViewer/PresidioDetectedEntities"; const BedrockPath = "@/components/view_logs/GuardrailViewer/BedrockGuardrailDetails"; +const skippedPreCall: Partial = { + guardrail_status: "not_run", + guardrail_mode: "pre_call", + guardrail_response: "no scannable content after message scoping", + start_time: null, + end_time: null, + duration: null, +}; + +const ranPostCall: Partial = { + guardrail_name: "ran-rail", + guardrail_status: "success", + guardrail_mode: "post_call", + start_time: 1_700_000_000, + end_time: 1_700_000_000.25, + duration: 0.25, +}; + describe("GuardrailViewer", () => { beforeEach(() => { vi.resetModules(); @@ -51,14 +70,7 @@ describe("GuardrailViewer", () => { it("renders not_run as NOT RUN (muted) and keeps it out of the evaluated and passed counts", async () => { const user = userEvent.setup(); - const data = makeGuardrailInformation({ - guardrail_status: "not_run", - guardrail_mode: "pre_call", - guardrail_response: "no scannable content after message scoping", - start_time: null, - end_time: null, - duration: null, - }); + const data = makeGuardrailInformation(skippedPreCall); renderWithProviders(); expect(screen.getByText(/0 guardrails evaluated/)).toBeInTheDocument(); @@ -74,22 +86,8 @@ describe("GuardrailViewer", () => { }); it("anchors the lifecycle timeline on timed entries when an untimed not_run entry sorts first", () => { - const skipped = makeGuardrailInformation({ - guardrail_name: "skipped-rail", - guardrail_status: "not_run", - guardrail_mode: "pre_call", - start_time: null, - end_time: null, - duration: null, - }); - const ran = makeGuardrailInformation({ - guardrail_name: "ran-rail", - guardrail_status: "success", - guardrail_mode: "post_call", - start_time: 1_700_000_000, - end_time: 1_700_000_000.25, - duration: 0.25, - }); + const skipped = makeGuardrailInformation({ ...skippedPreCall, guardrail_name: "skipped-rail" }); + const ran = makeGuardrailInformation(ranPostCall); renderWithProviders(); expect(screen.getByText(/1 guardrail evaluated/)).toBeInTheDocument(); From c3f52fe0d5cd646d2d9fd928830a17dd3ca9b46e Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 20:37:09 +0000 Subject: [PATCH 051/107] fix(guardrails): key not_run index rows by the sibling evaluation's guardrail_id A not_run entry from the base guardrail only carries guardrail_name, while the content filter's evaluated entry carries guardrail_id. Keyed apart, one request listed twice in the monitor for a logging_only guardrail (Not run and Passed). Resolve the id from a same-name sibling in the payload so the severity pick applies Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/guardrails/usage_tracking.py | 13 +++++++-- .../proxy/guardrails/test_usage_tracking.py | 28 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 6d72d94718f..3dfa4a96642 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -365,8 +365,17 @@ async def process_spend_logs_guardrail_usage( continue date_key = _date_str(start_time) - for entry in _parse_guardrail_info_from_payload(payload): - guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or "" + entries = _parse_guardrail_info_from_payload(payload) + ids_by_name = MappingProxyType( + { + e["guardrail_name"]: e["guardrail_id"] + for e in entries + if e.get("guardrail_id") and e.get("guardrail_name") + } + ) + for entry in entries: + guardrail_name = entry.get("guardrail_name") or "" + guardrail_id = entry.get("guardrail_id") or ids_by_name.get(guardrail_name) or guardrail_name if not guardrail_id: continue action = guardrail_status_to_action(entry.get("guardrail_status")) diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 0226e8b8e57..3aede9ba5e5 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -369,6 +369,34 @@ async def test_not_run_entries_are_indexed_but_not_counted_as_evaluations(): assert sorted(row["request_id"] for row in index_rows) == ["r1", "r2"] +@pytest.mark.asyncio +async def test_not_run_entry_shares_index_key_with_evaluated_sibling_of_same_name(): + """ + The not_run entry from the shared base guardrail carries only guardrail_name, + while the evaluated entry from the same guardrail (e.g. content filter on the + output of a logging_only run) carries its guardrail_id. Keying them differently + lists one request twice in the monitor, once as not_run and once as passed. + """ + prisma = _prisma() + payload = _payload("r1") + payload["metadata"] = json.dumps( + { + "guardrail_information": [ + {"guardrail_name": "cf", "guardrail_status": "not_run"}, + {"guardrail_name": "cf", "guardrail_id": "cf-uuid", "guardrail_status": "success"}, + {"guardrail_name": "other", "guardrail_status": "not_run"}, + ] + } + ) + + await process_spend_logs_guardrail_usage(prisma, [payload]) + + index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"] + assert sorted(row["guardrail_id"] for row in index_rows) == ["cf-uuid", "cf-uuid", "other"] + metrics_create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"] + assert (metrics_create["guardrail_id"], metrics_create["requests_evaluated"]) == ("cf-uuid", 1) + + @pytest.mark.asyncio async def test_batch_of_only_not_run_entries_writes_no_metrics_row(): prisma = _prisma() From 7ebb169a4d78f5af7e38a9b8c6dbeedce25f22a5 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 20:45:47 +0000 Subject: [PATCH 052/107] fix(guardrails): coalesce usage index rows per request and guardrail, keeping policy linkage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/guardrails/usage_tracking.py | 9 +++++---- .../proxy/guardrails/test_usage_tracking.py | 12 ++++++++++-- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 3dfa4a96642..4c8222f52ab 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -356,7 +356,7 @@ async def process_spend_logs_guardrail_usage( "flagged_count": 0, } ) - index_rows: Final[list[dict[str, object]]] = [] + index_rows_by_key: Final[dict[tuple[object, object], dict[str, object]]] = {} for payload in logs_to_process: request_id = payload.get("request_id") @@ -389,14 +389,15 @@ async def process_spend_logs_guardrail_usage( else: daily_guardrail[key]["flagged_count"] += 1 policy_id = entry.get("policy_id") - index_rows.append( - { + prior = index_rows_by_key.get((request_id, guardrail_id)) + if prior is None or (prior["policy_id"] is None and policy_id is not None): + index_rows_by_key[(request_id, guardrail_id)] = { "request_id": request_id, "guardrail_id": guardrail_id, "policy_id": policy_id, "start_time": start_time, } - ) + index_rows: Final = tuple(index_rows_by_key.values()) async with pending.lock: pending_metrics: Final = pending.metrics diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 3aede9ba5e5..58479d3d740 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -383,7 +383,12 @@ async def test_not_run_entry_shares_index_key_with_evaluated_sibling_of_same_nam { "guardrail_information": [ {"guardrail_name": "cf", "guardrail_status": "not_run"}, - {"guardrail_name": "cf", "guardrail_id": "cf-uuid", "guardrail_status": "success"}, + { + "guardrail_name": "cf", + "guardrail_id": "cf-uuid", + "policy_id": "pol-1", + "guardrail_status": "success", + }, {"guardrail_name": "other", "guardrail_status": "not_run"}, ] } @@ -392,7 +397,10 @@ async def test_not_run_entry_shares_index_key_with_evaluated_sibling_of_same_nam await process_spend_logs_guardrail_usage(prisma, [payload]) index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"] - assert sorted(row["guardrail_id"] for row in index_rows) == ["cf-uuid", "cf-uuid", "other"] + assert sorted((row["guardrail_id"], row["policy_id"]) for row in index_rows) == [ + ("cf-uuid", "pol-1"), + ("other", None), + ] metrics_create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"] assert (metrics_create["guardrail_id"], metrics_create["requests_evaluated"]) == ("cf-uuid", 1) From 586d51f15e92128c18fe7024bec6277a74831cb3 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 20:50:41 +0000 Subject: [PATCH 053/107] fix(guardrails): skip malformed guardrail entries instead of failing the usage batch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/guardrails/usage_tracking.py | 11 +++++----- .../proxy/guardrails/test_usage_tracking.py | 21 +++++++++++++++++++ 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 4c8222f52ab..b2ba6ec9ec5 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -356,12 +356,12 @@ async def process_spend_logs_guardrail_usage( "flagged_count": 0, } ) - index_rows_by_key: Final[dict[tuple[object, object], dict[str, object]]] = {} + index_rows_by_key: Final[dict[tuple[str, str], dict[str, object]]] = {} for payload in logs_to_process: request_id = payload.get("request_id") start_time = _parse_payload_start_time(payload) - if not request_id or start_time is None: + if not isinstance(request_id, str) or not request_id or start_time is None: continue date_key = _date_str(start_time) @@ -370,13 +370,14 @@ async def process_spend_logs_guardrail_usage( { e["guardrail_name"]: e["guardrail_id"] for e in entries - if e.get("guardrail_id") and e.get("guardrail_name") + if e.get("guardrail_id") and isinstance(e.get("guardrail_name"), str) } ) for entry in entries: - guardrail_name = entry.get("guardrail_name") or "" + raw_name = entry.get("guardrail_name") + guardrail_name = raw_name if isinstance(raw_name, str) else "" guardrail_id = entry.get("guardrail_id") or ids_by_name.get(guardrail_name) or guardrail_name - if not guardrail_id: + if not isinstance(guardrail_id, str) or not guardrail_id: continue action = guardrail_status_to_action(entry.get("guardrail_status")) if action != "not_run": diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 58479d3d740..56a8bf2f0b0 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -405,6 +405,27 @@ async def test_not_run_entry_shares_index_key_with_evaluated_sibling_of_same_nam assert (metrics_create["guardrail_id"], metrics_create["requests_evaluated"]) == ("cf-uuid", 1) +@pytest.mark.asyncio +async def test_malformed_not_run_entry_does_not_drop_the_batch(): + prisma = _prisma() + payload = _payload("r1") + payload["metadata"] = json.dumps( + { + "guardrail_information": [ + {"guardrail_name": ["not", "a", "string"], "guardrail_status": "not_run"}, + {"guardrail_name": "cf", "guardrail_id": "cf-uuid", "guardrail_status": "success"}, + ] + } + ) + + await process_spend_logs_guardrail_usage(prisma, [payload]) + + index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"] + assert [row["guardrail_id"] for row in index_rows] == ["cf-uuid"] + metrics_create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"] + assert metrics_create["requests_evaluated"] == 1 + + @pytest.mark.asyncio async def test_batch_of_only_not_run_entries_writes_no_metrics_row(): prisma = _prisma() From 937179bde1699b2c758a31f8850181d22fcb2302 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 20:53:11 +0000 Subject: [PATCH 054/107] fix(guardrails): never map empty guardrail names onto a sibling id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/guardrails/usage_tracking.py | 2 +- tests/test_litellm/proxy/guardrails/test_usage_tracking.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index b2ba6ec9ec5..7e11b69108b 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -370,7 +370,7 @@ async def process_spend_logs_guardrail_usage( { e["guardrail_name"]: e["guardrail_id"] for e in entries - if e.get("guardrail_id") and isinstance(e.get("guardrail_name"), str) + if e.get("guardrail_id") and isinstance(e.get("guardrail_name"), str) and e["guardrail_name"] } ) for entry in entries: diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 56a8bf2f0b0..13a53efcb27 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -412,8 +412,9 @@ async def test_malformed_not_run_entry_does_not_drop_the_batch(): payload["metadata"] = json.dumps( { "guardrail_information": [ - {"guardrail_name": ["not", "a", "string"], "guardrail_status": "not_run"}, - {"guardrail_name": "cf", "guardrail_id": "cf-uuid", "guardrail_status": "success"}, + {"guardrail_name": ["not", "a", "string"], "guardrail_status": "success"}, + {"guardrail_name": "", "guardrail_id": "cf-uuid", "guardrail_status": "success"}, + {"guardrail_status": "not_run"}, ] } ) @@ -423,7 +424,7 @@ async def test_malformed_not_run_entry_does_not_drop_the_batch(): index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"] assert [row["guardrail_id"] for row in index_rows] == ["cf-uuid"] metrics_create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"] - assert metrics_create["requests_evaluated"] == 1 + assert (metrics_create["guardrail_id"], metrics_create["requests_evaluated"]) == ("cf-uuid", 1) @pytest.mark.asyncio From 33fb6625ade5d3ee84d4138375272e35ce9d0133 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 21:18:00 +0000 Subject: [PATCH 055/107] test(guardrails): cover nameless evaluated entries in the malformed usage batch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/guardrails/test_usage_tracking.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 13a53efcb27..69ec098b840 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -414,7 +414,7 @@ async def test_malformed_not_run_entry_does_not_drop_the_batch(): "guardrail_information": [ {"guardrail_name": ["not", "a", "string"], "guardrail_status": "success"}, {"guardrail_name": "", "guardrail_id": "cf-uuid", "guardrail_status": "success"}, - {"guardrail_status": "not_run"}, + {"guardrail_status": "success"}, ] } ) From b37ce94075124b4429996cd52313d100fbb5212e Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 22:15:53 +0000 Subject: [PATCH 056/107] refactor(guardrails): rename scoped-out evaluation status from not_run to skipped The per-guardrail status a scoped-out evaluation records is now skipped, matching the skip_*_in_guardrail settings that cause it. Request-level rollup still maps it to not_run so the StandardLoggingPayload status contract is unchanged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 1 + .../chat/guardrail_translation/handler.py | 2 +- litellm/proxy/compliance_checks.py | 2 +- litellm/proxy/guardrails/usage_endpoints.py | 4 ++-- litellm/proxy/guardrails/usage_tracking.py | 8 +++---- litellm/types/utils.py | 3 ++- .../test_litellm_logging.py | 12 ++++++++++ .../test_openai_guardrail_handler.py | 8 +++---- .../proxy/guardrails/test_usage_endpoints.py | 10 ++++---- .../proxy/guardrails/test_usage_tracking.py | 22 ++++++++--------- .../test_compliance_endpoints.py | 12 +++++----- .../GuardrailsMonitor/LogViewer.test.tsx | 8 +++---- .../GuardrailsMonitor/LogViewer.tsx | 6 ++--- .../components/GuardrailsMonitor/mockData.ts | 2 +- .../GuardrailViewer/GuardrailViewer.test.tsx | 10 ++++---- .../GuardrailViewer/GuardrailViewer.tsx | 24 +++++++++---------- .../LogDetailContent.integration.test.tsx | 14 +++++------ .../LogDetailsDrawer/LogDetailContent.tsx | 12 +++++----- 18 files changed, 87 insertions(+), 73 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9ba9fd082f3..e88de4acd0e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6021,6 +6021,7 @@ def _get_status_fields( "failure": "guardrail_failed_to_respond", # legacy "guardrail_failed_to_respond": "guardrail_failed_to_respond", # direct "not_run": "not_run", + "skipped": "not_run", } # Set LLM API status diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 0f5096d0108..dca90e07421 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -214,7 +214,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response="no scannable content after message scoping", request_data=data, - guardrail_status="not_run", + guardrail_status="skipped", ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/compliance_checks.py b/litellm/proxy/compliance_checks.py index d9cc1d0f4fc..ef2d8fb6e20 100644 --- a/litellm/proxy/compliance_checks.py +++ b/litellm/proxy/compliance_checks.py @@ -26,7 +26,7 @@ class ComplianceChecker: def __init__(self, data: ComplianceCheckRequest): self.data = data - self.guardrails = [g for g in (data.guardrail_information or []) if g.get("guardrail_status") != "not_run"] + self.guardrails = [g for g in (data.guardrail_information or []) if g.get("guardrail_status") != "skipped"] def _get_guardrails_by_mode(self, mode: str) -> list[dict]: """ diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 556b6a4e919..651c4bb1963 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -42,7 +42,7 @@ if TYPE_CHECKING: router: Final = APIRouter() _EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({}) -_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"not_run": 0, "passed": 1, "flagged": 2, "blocked": 3}) +_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"skipped": 0, "passed": 1, "flagged": 2, "blocked": 3}) _T = TypeVar("_T") @@ -325,7 +325,7 @@ class UsageDetailResponse(BaseModel): class UsageLogEntry(BaseModel): id: str timestamp: str - action: str # blocked | passed | flagged | not_run + action: str # blocked | passed | flagged | skipped score: float | None latency_ms: float | None model: str | None diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 7e11b69108b..8a131fbfcde 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -193,12 +193,12 @@ async def _upsert_rows_with_retry( def guardrail_status_to_action(status: str | None) -> str: - """Map StandardLogging guardrail_status to blocked/passed/flagged/not_run.""" + """Map StandardLogging guardrail_status to blocked/passed/flagged/skipped.""" if not status: return "passed" s: Final = (status or "").lower() - if s == "not_run": - return "not_run" + if s == "skipped": + return "skipped" if "intervened" in s or "block" in s: return "blocked" if "flagged" in s or "fail" in s or "error" in s: @@ -380,7 +380,7 @@ async def process_spend_logs_guardrail_usage( if not isinstance(guardrail_id, str) or not guardrail_id: continue action = guardrail_status_to_action(entry.get("guardrail_status")) - if action != "not_run": + if action != "skipped": key = _MetricsKey(guardrail_id, date_key) daily_guardrail[key]["requests_evaluated"] += 1 if action == "passed": diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 1d73542c9bb..cf302b3f27f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3120,7 +3120,7 @@ class GuardrailMode(TypedDict, total=False): GuardrailStatus = Literal[ - "success", "guardrail_flagged", "guardrail_intervened", "guardrail_failed_to_respond", "not_run" + "success", "guardrail_flagged", "guardrail_intervened", "guardrail_failed_to_respond", "not_run", "skipped" ] # Fields on a guardrail record whose values can quote the caller's prompt: the payload sent to the @@ -3367,6 +3367,7 @@ class StandardLoggingPayloadStatusFields(TypedDict, total=False): - 'guardrail_intervened': Guardrail blocked or modified content - 'guardrail_failed_to_respond': Guardrail had technical failure - 'not_run': No guardrail was run + - 'skipped': Only used per guardrail entry, message scoping left the guardrail nothing to scan """ diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 70f9bae283b..6e67781944e 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -6926,6 +6926,18 @@ def test_get_status_fields_ranks_guardrail_flagged_between_success_and_intervene )["guardrail_status"] == "guardrail_intervened" +def test_get_status_fields_rolls_skipped_entries_up_to_not_run(): + """LIT-6314: a guardrail that message scoping left nothing to scan records a + skipped entry. At request level that means no guardrail ran, and a skipped + entry must never outrank a sibling that did evaluate.""" + skipped = {"guardrail_status": "skipped"} + + assert _get_status_fields("success", [skipped], None)["guardrail_status"] == "not_run" + assert _get_status_fields( + "success", [skipped, {"guardrail_status": "success"}], None + )["guardrail_status"] == "success" + + def test_get_error_information_redacts_provider_key_from_upstream_url(): """A pass-through upstream failure logs the httpx traceback, whose message quotes the upstream URL with the provider key in its query string. That diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 4f1163ed806..c6d01db45de 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1893,7 +1893,7 @@ class TestScanOnlyToolResults: assert data["messages"][4]["content"] == "and then?" -class TestNoScannableContentRecordsNotRun: +class TestNoScannableContentRecordsSkipped: """LIT-6314: a guardrail whose scoping leaves nothing to scan must still persist an evaluation record""" def _system_only_data(self) -> dict: @@ -1904,7 +1904,7 @@ class TestNoScannableContentRecordsNotRun: return metadata.get("standard_logging_guardrail_information") or [] @pytest.mark.asyncio - async def test_skipped_scan_records_not_run_entry(self): + async def test_skipped_scan_records_skipped_entry(self): handler = OpenAIChatCompletionsHandler() guardrail = MockGuardrail(guardrail_name="skip-system-guardrail") guardrail.skip_system_message_in_guardrail = True @@ -1916,7 +1916,7 @@ class TestNoScannableContentRecordsNotRun: entries = self._recorded_entries(data) assert len(entries) == 1 assert entries[0]["guardrail_name"] == "skip-system-guardrail" - assert entries[0]["guardrail_status"] == "not_run" + assert entries[0]["guardrail_status"] == "skipped" @pytest.mark.asyncio async def test_self_recording_guardrail_is_left_alone(self): @@ -1940,7 +1940,7 @@ class TestNoScannableContentRecordsNotRun: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) assert guardrail.last_inputs is not None - assert all(e.get("guardrail_status") != "not_run" for e in self._recorded_entries(data)) + assert all(e.get("guardrail_status") != "skipped" for e in self._recorded_entries(data)) class TestBuildBlockSseChunks: diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index db87e12ac88..1df0c1477e2 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -685,7 +685,7 @@ async def test_detail_prev_trend_query_is_bounded(): @pytest.mark.asyncio -async def test_logs_report_not_run_entries_as_not_run_not_passed(): +async def test_logs_report_skipped_entries_as_skipped_not_passed(): """LIT-6314: a guardrail that never scanned must not be reported as a pass in the drill-down.""" index_row = MagicMock() index_row.request_id = "req-nr" @@ -697,7 +697,7 @@ async def test_logs_report_not_run_entries_as_not_run_not_passed(): spend_log.startTime = datetime(2026, 4, 22) spend_log.metadata = { "guardrail_information": [ - {"guardrail_name": "db-1", "guardrail_status": "not_run", "duration": 0.0}, + {"guardrail_name": "db-1", "guardrail_status": "skipped", "duration": 0.0}, ] } prisma = _prisma(find_unique=_db_row(), index_find_many=[index_row]) @@ -715,11 +715,11 @@ async def test_logs_report_not_run_entries_as_not_run_not_passed(): end_date=END, user_api_key_dict=ADMIN, ) - assert [log.action for log in resp.logs] == ["not_run"] + assert [log.action for log in resp.logs] == ["skipped"] @pytest.mark.asyncio -async def test_logs_action_passed_filter_excludes_not_run_entries(): +async def test_logs_action_passed_filter_excludes_skipped_entries(): """LIT-6314: filtering the drill-down for passes must not return unscanned requests.""" index_row = MagicMock() index_row.request_id = "req-nr" @@ -729,7 +729,7 @@ async def test_logs_action_passed_filter_excludes_not_run_entries(): spend_log.request_id = "req-nr" spend_log.model = "gpt-4o-mini" spend_log.startTime = datetime(2026, 4, 22) - spend_log.metadata = {"guardrail_information": [{"guardrail_name": "db-1", "guardrail_status": "not_run"}]} + spend_log.metadata = {"guardrail_information": [{"guardrail_name": "db-1", "guardrail_status": "skipped"}]} prisma = _prisma(find_unique=_db_row(), index_find_many=[index_row]) prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[spend_log]) handler = _config_handler() diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index 69ec098b840..cb85eb8b5ec 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -350,15 +350,15 @@ async def test_zero_and_non_int_usage_counters_are_skipped(): @pytest.mark.asyncio -async def test_not_run_entries_are_indexed_but_not_counted_as_evaluations(): +async def test_skipped_entries_are_indexed_but_not_counted_as_evaluations(): """ - LIT-6314 records a not_run entry when message scoping leaves a guardrail + LIT-6314 records a skipped entry when message scoping leaves a guardrail nothing to scan. The guardrail never evaluated the request, so counting it as a passed evaluation would inflate daily pass rates; it still gets an index row so per-request drill-down finds the spend log. """ prisma = _prisma() - logs = [_payload("r1", guardrail_status="not_run"), _payload("r2")] + logs = [_payload("r1", guardrail_status="skipped"), _payload("r2")] await process_spend_logs_guardrail_usage(prisma, logs) @@ -370,26 +370,26 @@ async def test_not_run_entries_are_indexed_but_not_counted_as_evaluations(): @pytest.mark.asyncio -async def test_not_run_entry_shares_index_key_with_evaluated_sibling_of_same_name(): +async def test_skipped_entry_shares_index_key_with_evaluated_sibling_of_same_name(): """ - The not_run entry from the shared base guardrail carries only guardrail_name, + The skipped entry from the shared base guardrail carries only guardrail_name, while the evaluated entry from the same guardrail (e.g. content filter on the output of a logging_only run) carries its guardrail_id. Keying them differently - lists one request twice in the monitor, once as not_run and once as passed. + lists one request twice in the monitor, once as skipped and once as passed. """ prisma = _prisma() payload = _payload("r1") payload["metadata"] = json.dumps( { "guardrail_information": [ - {"guardrail_name": "cf", "guardrail_status": "not_run"}, + {"guardrail_name": "cf", "guardrail_status": "skipped"}, { "guardrail_name": "cf", "guardrail_id": "cf-uuid", "policy_id": "pol-1", "guardrail_status": "success", }, - {"guardrail_name": "other", "guardrail_status": "not_run"}, + {"guardrail_name": "other", "guardrail_status": "skipped"}, ] } ) @@ -406,7 +406,7 @@ async def test_not_run_entry_shares_index_key_with_evaluated_sibling_of_same_nam @pytest.mark.asyncio -async def test_malformed_not_run_entry_does_not_drop_the_batch(): +async def test_malformed_skipped_entry_does_not_drop_the_batch(): prisma = _prisma() payload = _payload("r1") payload["metadata"] = json.dumps( @@ -428,10 +428,10 @@ async def test_malformed_not_run_entry_does_not_drop_the_batch(): @pytest.mark.asyncio -async def test_batch_of_only_not_run_entries_writes_no_metrics_row(): +async def test_batch_of_only_skipped_entries_writes_no_metrics_row(): prisma = _prisma() - await process_spend_logs_guardrail_usage(prisma, [_payload("r1", guardrail_status="not_run")]) + await process_spend_logs_guardrail_usage(prisma, [_payload("r1", guardrail_status="skipped")]) assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 0 index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py index 8382a5ada96..9587e3d95ee 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py @@ -591,17 +591,17 @@ class TestModeMatching: assert mode in _guaranteed_modes(g_mode), (g_mode, mode) -class TestNotRunGuardrails: - """LIT-6314 logs a not_run entry for a guardrail that message scoping left nothing to scan.""" +class TestSkippedGuardrails: + """LIT-6314 logs a skipped entry for a guardrail that message scoping left nothing to scan.""" - def test_not_run_alone_never_evidences_compliance(self): + def test_skipped_alone_never_evidences_compliance(self): data = ComplianceCheckRequest( request_id="req-601", user_id="user-1", model="gpt-4", timestamp="2026-02-17T00:00:00Z", guardrail_information=[ - {"guardrail_name": "pii_detection", "guardrail_status": "not_run", "guardrail_mode": "pre_call"}, + {"guardrail_name": "pii_detection", "guardrail_status": "skipped", "guardrail_mode": "pre_call"}, ], ) results = {c.check_name: c.passed for c in ComplianceChecker(data).check_eu_ai_act()} @@ -609,7 +609,7 @@ class TestNotRunGuardrails: assert results["Content screened before LLM"] is False assert results["Audit record complete"] is False - def test_not_run_sibling_does_not_fail_a_passing_request(self): + def test_skipped_sibling_does_not_fail_a_passing_request(self): data = ComplianceCheckRequest( request_id="req-602", user_id="user-1", @@ -618,7 +618,7 @@ class TestNotRunGuardrails: pii_detected=True, guardrail_information=[ {"guardrail_name": "pii_detection", "guardrail_status": "success", "guardrail_mode": "pre_call"}, - {"guardrail_name": "system_only", "guardrail_status": "not_run", "guardrail_mode": "pre_call"}, + {"guardrail_name": "system_only", "guardrail_status": "skipped", "guardrail_mode": "pre_call"}, ], ) results = {c.check_name: c.passed for c in ComplianceChecker(data).check_gdpr()} diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx index 083b1e5f3e2..af5284830dc 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx @@ -96,14 +96,14 @@ describe("GuardrailsMonitor LogViewer drawer", () => { }); }); -describe("GuardrailsMonitor LogViewer not_run rows", () => { - it("renders a not_run log as a neutral Not run badge instead of a pass or failure", () => { +describe("GuardrailsMonitor LogViewer skipped rows", () => { + it("renders a skipped log as a neutral Skipped badge instead of a pass or failure", () => { renderWithProviders( - , + , ); const row = screen.getByRole("button", { name: /system prompt only/ }); - expect(within(row).getByText("Not run")).toHaveClass("text-muted-foreground"); + expect(within(row).getByText("Skipped")).toHaveClass("text-muted-foreground"); expect(within(row).queryByText("Passed")).not.toBeInTheDocument(); expect(within(row).queryByText("Blocked")).not.toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx index 2abd699ba86..b113c49e0a3 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx @@ -10,15 +10,15 @@ import type { LogEntry as ViewLogsLogEntry } from "@/components/view_logs/column import type { LogEntry } from "./mockData"; const actionConfig: Record< - "blocked" | "passed" | "flagged" | "not_run", + "blocked" | "passed" | "flagged" | "skipped", { icon: React.ElementType; color: string; bg: string; border: string; label: string } > = { - not_run: { + skipped: { icon: MinusCircle, color: "text-muted-foreground", bg: "bg-muted", border: "border-border", - label: "Not run", + label: "Skipped", }, blocked: { icon: X, diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts b/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts index 591d5cd3edd..7053efcf88d 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts @@ -10,7 +10,7 @@ export interface LogEntry { input_snippet?: string; output_snippet?: string; score?: number; - action: "blocked" | "passed" | "flagged" | "not_run"; + action: "blocked" | "passed" | "flagged" | "skipped"; model?: string; reason?: string; latency_ms?: number; diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx index 7f343211596..740ef857690 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx @@ -16,7 +16,7 @@ const PresidioPath = "@/components/view_logs/GuardrailViewer/PresidioDetectedEnt const BedrockPath = "@/components/view_logs/GuardrailViewer/BedrockGuardrailDetails"; const skippedPreCall: Partial = { - guardrail_status: "not_run", + guardrail_status: "skipped", guardrail_mode: "pre_call", guardrail_response: "no scannable content after message scoping", start_time: null, @@ -68,15 +68,15 @@ describe("GuardrailViewer", () => { expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); }); - it("renders not_run as NOT RUN (muted) and keeps it out of the evaluated and passed counts", async () => { + it("renders skipped as SKIPPED (muted) and keeps it out of the evaluated and passed counts", async () => { const user = userEvent.setup(); const data = makeGuardrailInformation(skippedPreCall); renderWithProviders(); expect(screen.getByText(/0 guardrails evaluated/)).toBeInTheDocument(); expect(screen.getByText(/0 Passed/)).toHaveClass("text-muted-foreground"); - expect(screen.getByText(/1 Not run/)).toBeInTheDocument(); - const badge = screen.getByText("NOT RUN"); + expect(screen.getByText(/1 Skipped/)).toBeInTheDocument(); + const badge = screen.getByText("SKIPPED"); expect(badge).toHaveClass("text-muted-foreground"); expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); expect(screen.queryByText(/^T\+/)).not.toBeInTheDocument(); @@ -85,7 +85,7 @@ describe("GuardrailViewer", () => { expect(screen.getByText("no scannable content after message scoping")).toBeInTheDocument(); }); - it("anchors the lifecycle timeline on timed entries when an untimed not_run entry sorts first", () => { + it("anchors the lifecycle timeline on timed entries when an untimed skipped entry sorts first", () => { const skipped = makeGuardrailInformation({ ...skippedPreCall, guardrail_name: "skipped-rail" }); const ran = makeGuardrailInformation(ranPostCall); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx index 1de0e3878b2..c67d682a233 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx @@ -134,13 +134,13 @@ const getTotalMasked = (entry: GuardrailInformation): number => { ); }; -type EntryOutcome = "passed" | "flagged" | "failed" | "not_run"; +type EntryOutcome = "passed" | "flagged" | "failed" | "skipped"; const getEntryOutcome = (entry: GuardrailInformation): EntryOutcome => { const status = (entry.guardrail_status ?? "").toLowerCase(); if (status === "success") return "passed"; if (status === "guardrail_flagged") return "flagged"; - if (status === "not_run") return "not_run"; + if (status === "skipped") return "skipped"; return "failed"; }; @@ -150,18 +150,18 @@ const OUTCOME_LABEL: Record = { passed: "PASSED", flagged: "FLAGGED", failed: "FAILED", - not_run: "NOT RUN", + skipped: "SKIPPED", }; const OUTCOME_BADGE_CLASS: Record = { passed: "bg-success/15 text-success border border-success/20", flagged: "bg-warning/15 text-warning border border-warning/20", failed: "bg-destructive/15 text-destructive border border-destructive/20", - not_run: "bg-muted text-muted-foreground border border-border", + skipped: "bg-muted text-muted-foreground border border-border", }; const getHeaderOutcome = (counts: { evaluated: number; passed: number; flagged: number }): EntryOutcome => { - if (counts.evaluated === 0) return "not_run"; + if (counts.evaluated === 0) return "skipped"; if (counts.passed === counts.evaluated) return "passed"; if (counts.passed + counts.flagged === counts.evaluated) return "flagged"; return "failed"; @@ -242,7 +242,7 @@ const FlagCircleIcon = ({ className }: { className?: string }) => ( const OutcomeIcon = ({ outcome }: { outcome: EntryOutcome }) => { if (outcome === "passed") return ; if (outcome === "flagged") return ; - if (outcome === "not_run") return ; + if (outcome === "skipped") return ; return ; }; @@ -675,7 +675,7 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => {
)} - {outcome === "not_run" && typeof guardrailResponse === "string" && ( + {outcome === "skipped" && typeof guardrailResponse === "string" && (

{guardrailResponse}

)} @@ -717,8 +717,8 @@ const GuardrailViewer = ({ data, accessToken, logEntry }: GuardrailViewerProps) const passedCount = guardrailEntries.filter(isEntrySuccess).length; const flaggedCount = guardrailEntries.filter((e) => getEntryOutcome(e) === "flagged").length; - const notRunCount = guardrailEntries.filter((e) => getEntryOutcome(e) === "not_run").length; - const evaluatedCount = guardrailEntries.length - notRunCount; + const skippedCount = guardrailEntries.filter((e) => getEntryOutcome(e) === "skipped").length; + const evaluatedCount = guardrailEntries.length - skippedCount; const allPassed = evaluatedCount > 0 && passedCount === evaluatedCount; const headerOutcome = getHeaderOutcome({ evaluated: evaluatedCount, passed: passedCount, flagged: flaggedCount }); @@ -778,11 +778,11 @@ const GuardrailViewer = ({ data, accessToken, logEntry }: GuardrailViewerProps) {flaggedCount} Flagged )} - {notRunCount > 0 && ( + {skippedCount > 0 && ( - {notRunCount} Not run + {skippedCount} Skipped )} diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx index 3637b6c55a3..cbc377c1b73 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx @@ -637,20 +637,20 @@ describe("GuardrailJumpLink", () => { }); it.each([ - [["success", "not_run"], "text-success", "\u2713"], - [["guardrail_intervened", "not_run"], "text-destructive", "\u2717"], - ])("ignores not_run when styling %j as %s", (statuses, expectedClass, glyph) => { + [["success", "skipped"], "text-success", "\u2713"], + [["guardrail_intervened", "skipped"], "text-destructive", "\u2717"], + ])("ignores skipped when styling %j as %s", (statuses, expectedClass, glyph) => { render( ({ guardrail_status: s }))} />); - const pill = screen.getByText(/1 guardrail evaluated, 1 not run/); + const pill = screen.getByText(/1 guardrail evaluated, 1 skipped/); expect(pill).toHaveClass(expectedClass); expect(pill).toHaveTextContent(glyph); }); - it("renders an all not_run request as neutral rather than passed", () => { - render(); + it("renders an all skipped request as neutral rather than passed", () => { + render(); - const pill = screen.getByText(/0 guardrails evaluated, 1 not run/); + const pill = screen.getByText(/0 guardrails evaluated, 1 skipped/); expect(pill).toHaveClass("text-muted-foreground"); expect(pill).not.toHaveTextContent("\u2713"); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx index 30cd96935f7..e052faa2228 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx @@ -700,15 +700,15 @@ const GUARDRAIL_JUMP_LINK_STYLE = { passed: { className: "border border-success/20 bg-success/10 text-success", glyph: "\u2713" }, flagged: { className: "border border-warning/20 bg-warning/10 text-warning", glyph: "\u26A0" }, failed: { className: "border border-destructive/20 bg-destructive/10 text-destructive", glyph: "\u2717" }, - not_run: { className: "border border-border bg-muted text-muted-foreground", glyph: "\u2013" }, + skipped: { className: "border border-border bg-muted text-muted-foreground", glyph: "\u2013" }, } as const; const isPassedStatus = (status: unknown) => status === "pass" || status === "passed" || status === "success"; const isFlaggedStatus = (status: unknown) => status === "flagged" || status === "guardrail_flagged"; -const isNotRunStatus = (status: unknown) => status === "not_run"; +const isSkippedStatus = (status: unknown) => status === "skipped"; const guardrailJumpLinkOutcome = (evaluated: unknown[]): keyof typeof GUARDRAIL_JUMP_LINK_STYLE => { - if (evaluated.length === 0) return "not_run"; + if (evaluated.length === 0) return "skipped"; if (evaluated.every(isPassedStatus)) return "passed"; if (evaluated.every((s) => isPassedStatus(s) || isFlaggedStatus(s))) return "flagged"; return "failed"; @@ -716,8 +716,8 @@ const guardrailJumpLinkOutcome = (evaluated: unknown[]): keyof typeof GUARDRAIL_ export function GuardrailJumpLink({ guardrailEntries }: { guardrailEntries: any[] }) { const statuses = guardrailEntries.map((e) => e?.guardrail_status || e?.status); - const evaluated = statuses.filter((s) => !isNotRunStatus(s)); - const notRunCount = statuses.length - evaluated.length; + const evaluated = statuses.filter((s) => !isSkippedStatus(s)); + const skippedCount = statuses.length - evaluated.length; const { className, glyph } = GUARDRAIL_JUMP_LINK_STYLE[guardrailJumpLinkOutcome(evaluated)]; const handleClick = () => { @@ -743,7 +743,7 @@ export function GuardrailJumpLink({ guardrailEntries }: { guardrailEntries: any[ > {glyph} {evaluated.length} guardrail {evaluated.length !== 1 ? "s" : ""} evaluated - {notRunCount > 0 ? `, ${notRunCount} not run` : ""} + {skippedCount > 0 ? `, ${skippedCount} skipped` : ""} {"\u2193"} From c788769129e23261d7be091cac9aa036c2c04b09 Mon Sep 17 00:00:00 2001 From: Matthew Howard <78384492+matthowardcohere@users.noreply.github.com> Date: Mon, 14 Sep 2026 15:18:14 -0700 Subject: [PATCH 057/107] fix(terraform): recognize a credential name conflict and adopt the existing credential The proxy answered a duplicate credential_name with a 500 carrying Prisma's unique-constraint message, so terraform apply against a credential the state had lost died with an opaque database error. Classify that response as a conflict, adopt the existing credential with a PATCH that carries model_id, and set the resource ID only once the adopt succeeds so a failed PATCH does not taint a credential this run never owned Squashed from the commits on #39745 with authorship preserved Fixes https://github.com/BerriAI/terraform-provider-litellm/issues/8 --- terraform/provider/CHANGELOG.md | 1 + .../litellm/resource_credential_crud.go | 26 +++ .../litellm/resource_credential_crud_test.go | 157 ++++++++++++++++++ terraform/provider/litellm/utils.go | 32 ++++ 4 files changed, 216 insertions(+) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 8c0ef5a8b15..3daee3250b6 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -38,6 +38,7 @@ longer signal it. ### Fixed - **key**: An update that changes `team_id` and fails because the key was already cascade-deleted along with its previous team now recovers by recreating the key under the new team, instead of aborting the apply. The key's absence is confirmed against the proxy first, so an unrelated failure still errors out, and a `team_id` change between two teams that both still exist stays a plain in-place update +- **credential**: `litellm_credential` create now adopts an existing credential on a `credential_name` conflict instead of failing with a 500; `apply` is idempotent again once state loses track of a credential that still exists on the proxy - **team**: Read now decodes the `team_info` envelope `/team/info` actually returns, so team attributes refresh from the proxy instead of always falling back to the prior state - **key**: Read now unwraps the `info` envelope `/key/info` actually returns; previously reads mapped nothing back into state, so drift on a key was never detected - **key**: Read now picks up `model_rpm_limit`, `model_tpm_limit`, `guardrails`, `tags`, `enforced_params`, `allowed_passthrough_routes`, `rpm_limit_type`, `tpm_limit_type` and `prompts` from `info.metadata`, where the proxy actually stores them; previously they stayed empty in state, so a matching config showed a permanent phantom diff on them and out-of-band changes to them were never detected diff --git a/terraform/provider/litellm/resource_credential_crud.go b/terraform/provider/litellm/resource_credential_crud.go index dd9aef64f76..cb5031ee04e 100644 --- a/terraform/provider/litellm/resource_credential_crud.go +++ b/terraform/provider/litellm/resource_credential_crud.go @@ -88,6 +88,26 @@ func resourceLiteLLMCredentialCreate(d *schema.ResourceData, m interface{}) erro err = handleCredentialAPIResponse(resp, nil, client) if err != nil { + // If a credential with this name already exists, adopt it instead of + // failing: take ownership and update the existing credential's + // values (merged onto whatever it already had - not a full replace) + // rather than erroring on the unique-constraint conflict. See + // https://github.com/BerriAI/terraform-provider-litellm/issues/8. + if err.Error() == "credential_conflict" { + log.Printf("[WARN] Credential %q already exists; adopting it and updating to match configuration.", credentialName) + d.SetId(credentialName) + if updateErr := resourceLiteLLMCredentialUpdate(d, m); updateErr != nil { + // Adoption failed before this run took ownership of + // anything real. Clear the ID so create is reported as + // failed outright (matching pre-adoption behavior) instead + // of tainting state for a credential this run doesn't own - + // state that would otherwise get destroyed on the next + // apply. + d.SetId("") + return fmt.Errorf("failed to adopt existing credential %q: %w", credentialName, updateErr) + } + return nil + } return fmt.Errorf("failed to create credential: %w", err) } @@ -142,6 +162,7 @@ func resourceLiteLLMCredentialUpdate(d *schema.ResourceData, m interface{}) erro client := m.(*Client) credentialName := d.Id() + modelID := d.Get("model_id").(string) credentialInfo := d.Get("credential_info").(map[string]interface{}) credentialValues := d.Get("credential_values").(map[string]interface{}) @@ -157,8 +178,13 @@ func resourceLiteLLMCredentialUpdate(d *schema.ResourceData, m interface{}) erro credValuesMap[k] = v } + // model_id must travel with the update the same way it does on create, + // so the proxy's model-based credential resolution still applies. Without + // it, updating (or adopting) a model_id-scoped credential silently loses + // that association. credentialRequest := CredentialRequest{ CredentialName: credentialName, + ModelID: modelID, CredentialInfo: credInfoMap, CredentialValues: credValuesMap, } diff --git a/terraform/provider/litellm/resource_credential_crud_test.go b/terraform/provider/litellm/resource_credential_crud_test.go index 3398e58dd13..6e0b818fe33 100644 --- a/terraform/provider/litellm/resource_credential_crud_test.go +++ b/terraform/provider/litellm/resource_credential_crud_test.go @@ -3,6 +3,7 @@ package litellm import ( "encoding/json" "fmt" + "io" "net/http" "net/http/httptest" "sync/atomic" @@ -199,3 +200,159 @@ func TestRetryCredentialRead_ConnectionError(t *testing.T) { // Connection error should not be retried (not a "credential_not_found") fmt.Printf("connection error (expected): %v\n", err) } + +// conflictServer builds the shared conflict-then-recover mock used by the +// adoption tests below. patchStatus/patchBody control the PATCH response, so +// callers can exercise both the success and failure paths. +func conflictServer(t *testing.T, patchStatus int, patchBody string) (*httptest.Server, *int32, *int32, *[]byte) { + t.Helper() + var createCalls, patchCalls int32 + var capturedPatchBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost && r.URL.Path == "/credentials": + atomic.AddInt32(&createCalls, 1) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error":{"message":"Unique constraint failed on the fields: (` + "`credential_name`" + `)","type":"internal_server_error","code":"500"}}`)) + case r.Method == http.MethodPatch: + atomic.AddInt32(&patchCalls, 1) + if r.URL.Path != "/credentials/conflict-test" { + t.Errorf("PATCH went to %q, want /credentials/conflict-test", r.URL.Path) + } + body, _ := io.ReadAll(r.Body) + capturedPatchBody = body + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(patchStatus) + w.Write([]byte(patchBody)) + case r.Method == http.MethodGet: + if r.URL.Path != "/credentials/by_name/conflict-test" || r.URL.Query().Get("model_id") != "model-1" { + t.Errorf("GET went to %q (query %q), want /credentials/by_name/conflict-test?model_id=model-1", r.URL.Path, r.URL.RawQuery) + } + resp := CredentialResponse{CredentialName: "conflict-test", CredentialInfo: map[string]interface{}{}} + body, _ := json.Marshal(resp) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write(body) + default: + http.NotFound(w, r) + } + })) + return srv, &createCalls, &patchCalls, &capturedPatchBody +} + +// A credential that already exists in LiteLLM (created out of band, or left +// behind by a prior apply that dropped state) must be adopted on create +// instead of failing on the credential_name unique-constraint conflict, and +// the adopt PATCH must carry model_id so model-based credential resolution +// still applies (previously dropped - see +// https://github.com/BerriAI/litellm/pull/39745). +func TestResourceLiteLLMCredentialCreate_AdoptsOnConflict(t *testing.T) { + srv, createCalls, patchCalls, patchBody := conflictServer(t, http.StatusOK, `{}`) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMCredential().Schema, map[string]interface{}{ + "credential_name": "conflict-test", + "model_id": "model-1", + "credential_info": map[string]interface{}{"custom_llm_provider": "bedrock"}, + "credential_values": map[string]interface{}{"aws_access_key_id": "val"}, + }) + + if err := resourceLiteLLMCredentialCreate(d, client); err != nil { + t.Fatalf("expected create to adopt the existing credential, got error: %v", err) + } + if d.Id() != "conflict-test" { + t.Fatalf("expected ID %q, got %q", "conflict-test", d.Id()) + } + if got := atomic.LoadInt32(createCalls); got != 1 { + t.Fatalf("expected exactly 1 POST /credentials call, got %d", got) + } + if got := atomic.LoadInt32(patchCalls); got != 1 { + t.Fatalf("expected the conflict to trigger exactly 1 PATCH (adopt-and-update), got %d", got) + } + + var sent map[string]interface{} + if err := json.Unmarshal(*patchBody, &sent); err != nil { + t.Fatalf("PATCH body was not valid JSON: %v (%s)", err, *patchBody) + } + if sent["credential_name"] != "conflict-test" { + t.Errorf("PATCH body credential_name = %v, want conflict-test", sent["credential_name"]) + } + if sent["model_id"] != "model-1" { + t.Errorf("PATCH body model_id = %v, want model-1 (adoption must not drop model-based credential resolution)", sent["model_id"]) + } + credInfo, _ := sent["credential_info"].(map[string]interface{}) + if credInfo["custom_llm_provider"] != "bedrock" { + t.Errorf("PATCH body credential_info = %v, want custom_llm_provider=bedrock", sent["credential_info"]) + } +} + +// If the adopt PATCH itself fails, create must not have set the resource ID +// for a credential this run doesn't own - otherwise Terraform taints the +// entry and the *next* apply destroys a credential nobody here created. +func TestResourceLiteLLMCredentialCreate_FailedAdoptDoesNotTaint(t *testing.T) { + srv, createCalls, patchCalls, _ := conflictServer(t, http.StatusInternalServerError, `{"error":{"message":"Internal Server Error"}}`) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMCredential().Schema, map[string]interface{}{ + "credential_name": "conflict-test", + "credential_info": map[string]interface{}{}, + "credential_values": map[string]interface{}{"key": "val"}, + }) + + err := resourceLiteLLMCredentialCreate(d, client) + if err == nil { + t.Fatal("expected an error when the adopt PATCH fails, got nil") + } + if got := atomic.LoadInt32(createCalls); got != 1 { + t.Fatalf("expected exactly 1 POST /credentials call, got %d", got) + } + if got := atomic.LoadInt32(patchCalls); got != 1 { + t.Fatalf("expected exactly 1 PATCH attempt, got %d", got) + } + if d.Id() != "" { + t.Fatalf("resource ID must stay empty after a failed adopt, got %q (a tainted entry would be destroyed on the next apply)", d.Id()) + } +} + +// A non-conflict failure (a plain 500, for example) must return the original +// error and never attempt to adopt anything. +func TestResourceLiteLLMCredentialCreate_NonConflictErrorDoesNotAdopt(t *testing.T) { + var createCalls, patchCalls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost && r.URL.Path == "/credentials": + atomic.AddInt32(&createCalls, 1) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error":{"message":"Internal Server Error","type":"internal_server_error"}}`)) + case r.Method == http.MethodPatch: + atomic.AddInt32(&patchCalls, 1) + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMCredential().Schema, map[string]interface{}{ + "credential_name": "some-cred", + "credential_info": map[string]interface{}{}, + "credential_values": map[string]interface{}{"key": "val"}, + }) + + err := resourceLiteLLMCredentialCreate(d, client) + if err == nil { + t.Fatal("expected an error for a non-conflict failure, got nil") + } + if got := atomic.LoadInt32(&patchCalls); got != 0 { + t.Fatalf("expected no PATCH attempt for a non-conflict error, got %d", got) + } + if d.Id() != "" { + t.Fatalf("resource ID must stay empty on a non-conflict failure, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/utils.go b/terraform/provider/litellm/utils.go index 5e81766d3f3..a123f5d350a 100644 --- a/terraform/provider/litellm/utils.go +++ b/terraform/provider/litellm/utils.go @@ -202,6 +202,35 @@ func isCredentialNotFoundError(errResp ErrorResponse) bool { return false } +// isCredentialConflictError checks if the error response indicates a credential +// name collision. LiteLLM surfaces this as a 500 carrying the underlying Prisma +// unique-constraint message on credential_name. See +// https://github.com/BerriAI/terraform-provider-litellm/issues/8. +func isCredentialConflictError(errResp ErrorResponse) bool { + if msg, ok := errResp.Error.Message.(string); ok { + if strings.Contains(msg, "Unique constraint failed") && strings.Contains(msg, "credential_name") { + return true + } + } + + if msgMap, ok := errResp.Error.Message.(map[string]interface{}); ok { + if errStr, ok := msgMap["error"].(string); ok { + if strings.Contains(errStr, "Unique constraint failed") && strings.Contains(errStr, "credential_name") { + return true + } + } + } + + // Check Detail.Error field for LiteLLM proxy error format + if errResp.Detail.Error != "" { + if strings.Contains(errResp.Detail.Error, "Unique constraint failed") && strings.Contains(errResp.Detail.Error, "credential_name") { + return true + } + } + + return false +} + // handleCredentialAPIResponse handles API responses specifically for credential operations func handleCredentialAPIResponse(resp *http.Response, result interface{}, client *Client) error { bodyBytes, err := io.ReadAll(resp.Body) @@ -219,6 +248,9 @@ func handleCredentialAPIResponse(resp *http.Response, result interface{}, client if isCredentialNotFoundError(errResp) { return fmt.Errorf("credential_not_found") } + if isCredentialConflictError(errResp) { + return fmt.Errorf("credential_conflict") + } } return fmt.Errorf("API request failed: Status: %s, Response: %s", resp.Status, client.redactSensitiveData(string(bodyBytes))) From 96332f75e857a2deda2d5db6db00deb79185db17 Mon Sep 17 00:00:00 2001 From: ryan Date: Mon, 14 Sep 2026 22:28:39 +0000 Subject: [PATCH 058/107] refactor(proxy): move bulk user creation to POST /management/v1/users/bulk Follows the Management API modernization design: plural resource under /management/v1, {data, meta} response envelope, unknown request fields rejected with 422, and RFC 9457 problem+json for request-level errors. Body validation failures under /management/v1 now answer 422 instead of the 400 query-parameter problem. /user/bulk_new is removed. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 2 +- litellm/proxy/auth/route_checks.py | 4 +- litellm/proxy/list_api/common.py | 27 +++ .../internal_user_endpoints.py | 68 ------- .../management_v1/__init__.py | 4 + .../management_v1/users.py | 105 ++++++++++ .../management_helpers/bulk_user_creation.py | 24 ++- litellm/proxy/proxy_server.py | 26 +-- .../internal_user_endpoints.py | 28 ++- .../endpointaudit/coverage_allowlist.txt | 1 - .../management_v1/test_users.py | 122 +++++++++++ .../test_bulk_user_creation.py | 61 +++--- .../proxy_server/test_exception_handlers.py | 32 ++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 192 +++++++++--------- 14 files changed, 462 insertions(+), 234 deletions(-) create mode 100644 litellm/proxy/management_endpoints/management_v1/users.py create mode 100644 tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ed182460644..48a324bff71 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -656,7 +656,7 @@ class LiteLLMRoutes(enum.Enum): [ # user "/user/new", - "/user/bulk_new", + "/management/v1/users/bulk", "/user/update", "/user/bulk_update", "/user/delete", diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 051ab13c058..730ab30f65f 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -24,7 +24,7 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset( [ # user "/user/new", - "/user/bulk_new", + "/management/v1/users/bulk", "/user/delete", "/user/bulk_update", # team @@ -756,7 +756,7 @@ class RouteChecks: _ADMIN_VIEWER_BLOCKED_WRITE_ROUTES = frozenset( [ "/user/new", - "/user/bulk_new", + "/management/v1/users/bulk", "/user/delete", "/user/bulk_update", "/team/new", diff --git a/litellm/proxy/list_api/common.py b/litellm/proxy/list_api/common.py index 7ef2827f30e..a0b0d45848f 100644 --- a/litellm/proxy/list_api/common.py +++ b/litellm/proxy/list_api/common.py @@ -1,5 +1,6 @@ """Contract machinery shared by every LiteLLM-defined list route, on any surface.""" +from collections.abc import Sequence from typing import Final from urllib.parse import urlencode @@ -7,6 +8,7 @@ from fastapi import Request from fastapi.dependencies.utils import get_flat_params from fastapi.params import ParamTypes from fastapi.responses import JSONResponse +from typing_extensions import ReadOnly, TypedDict from litellm.types.proxy.management_endpoints.management_v1 import ( ListLinks, @@ -56,6 +58,31 @@ def escape_like(value: str) -> str: return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") +class ValidationErrorDetail(TypedDict): + """The two keys of a pydantic/FastAPI validation error a problem document needs.""" + + loc: ReadOnly[tuple[int | str, ...]] + msg: ReadOnly[str] + + +def request_validation_problem(errors: Sequence[ValidationErrorDetail]) -> ProblemDetail: + """A body that fails validation (an unknown field included) is 422; a bad query parameter is 400.""" + detail: Final = "; ".join(f"{'.'.join(str(part) for part in error['loc'][1:])}: {error['msg']}" for error in errors) + if any(error["loc"] and error["loc"][0] == "body" for error in errors): + return ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}invalid-request-body", + title="Invalid request body", + status=422, + detail=detail or "The request body is invalid.", + ) + return ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}invalid-query-parameter", + title="Invalid query parameter", + status=400, + detail=detail or "The request query parameters are invalid.", + ) + + def unknown_query_param_problem(unknown: tuple[str, ...], allowed: tuple[str, ...]) -> ProblemDetail: return ProblemDetail( type=f"{PROBLEM_TYPE_BASE}unknown-query-parameter", diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 5f570485ed3..e3efda507f6 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -5,7 +5,6 @@ Internal User Management Endpoints These are members of a Team on LiteLLM /user/new -/user/bulk_new /user/update /user/bulk_update /user/delete @@ -78,8 +77,6 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( - BulkNewUserRequest, - BulkNewUserResponse, BulkUpdateUserRequest, BulkUpdateUserResponse, UserListResponse, @@ -640,71 +637,6 @@ async def new_user( raise handle_exception_on_proxy(e) -@router.post( - "/user/bulk_new", - tags=["Internal User management"], - dependencies=[Depends(user_api_key_auth)], - response_model=BulkNewUserResponse, -) -@management_endpoint_wrapper -async def bulk_new_user( - data: BulkNewUserRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection -) -> BulkNewUserResponse: - """ - Create up to 500 internal users in one request, optionally adding each one to teams. - - Every entry in `users` takes the same fields as `/user/new`, with two differences: `auto_create_key` - defaults to `false` (opt in per user to also get a virtual key back) and `send_invite_email` is not - supported. Rows are validated together (duplicate ids or emails, unknown teams, roles the caller may not - grant), inserted in one statement, and each referenced team is written once for all of its new members. - - Rows fail independently: a bad row is reported in `results` with `success: false` and an `error`, and the - other rows still get created. A user that was created but could not be added to one of its teams is - reported with `success: true`, `teams` listing where they did land, and `error` naming the failed team. - The whole request is rejected with 403 only if creating the valid rows would exceed the license seat limit. - - Usage Example - - ```shell - curl -X POST "http://localhost:4000/user/bulk_new" \\ - -H "Content-Type: application/json" \\ - -H "Authorization: Bearer sk-1234" \\ - -d '{ - "users": [ - {"user_email": "a@example.com", "user_role": "internal_user", "teams": ["team-1"]}, - {"user_email": "b@example.com", "user_role": "internal_user", "auto_create_key": true} - ] - }' - ``` - - Returns `results` (one entry per input row, in order, with `user_id`, `user_email`, `success`, `teams`, - `key`, `error`), `total_requested`, `successful_creations` and `failed_creations`. - """ - from litellm.proxy.management_helpers.bulk_user_creation import bulk_create_users - from litellm.proxy.proxy_server import ( - _license_check, # pyright: ignore[reportPrivateUsage] # same proxy license singleton /user/new reads - litellm_proxy_admin_name, - prisma_client, - user_api_key_cache, - ) - - if prisma_client is None: - raise HTTPException(status_code=400, detail=CommonProxyErrors.db_not_connected_error.value) - try: - return await bulk_create_users( - users=data.users, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - license_check=_license_check, - litellm_proxy_admin_name=litellm_proxy_admin_name, - user_api_key_cache=user_api_key_cache, - ) - except Exception as e: # noqa: BLE001 # normalize every failure to the proxy exception contract - verbose_proxy_logger.exception("/user/bulk_new: Exception occured") - raise handle_exception_on_proxy(e) - - @router.get( "/user/available_roles", tags=["Internal User management"], diff --git a/litellm/proxy/management_endpoints/management_v1/__init__.py b/litellm/proxy/management_endpoints/management_v1/__init__.py index 342e6525cda..a2172162dae 100644 --- a/litellm/proxy/management_endpoints/management_v1/__init__.py +++ b/litellm/proxy/management_endpoints/management_v1/__init__.py @@ -10,9 +10,13 @@ from litellm.proxy.management_endpoints.management_v1.budgets import ( from litellm.proxy.management_endpoints.management_v1.spend_logs import ( router as spend_logs_router, ) +from litellm.proxy.management_endpoints.management_v1.users import ( + router as users_router, +) router: Final = APIRouter() router.include_router(budgets_router) router.include_router(spend_logs_router) +router.include_router(users_router) __all__ = ["router"] diff --git a/litellm/proxy/management_endpoints/management_v1/users.py b/litellm/proxy/management_endpoints/management_v1/users.py new file mode 100644 index 00000000000..33a55c08513 --- /dev/null +++ b/litellm/proxy/management_endpoints/management_v1/users.py @@ -0,0 +1,105 @@ +"""`POST /management/v1/users/bulk`.""" + +from typing import Annotated, Final + +from fastapi import APIRouter, Depends + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE, ManagementProblem +from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX +from litellm.proxy.management_helpers.bulk_user_creation import bulk_create_users +from litellm.proxy.management_helpers.utils import ( + management_endpoint_wrapper, # pyright: ignore[reportUnknownVariableType] # legacy untyped decorator +) +from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( + BulkNewUserRequest, + BulkNewUserResponse, +) +from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail + +router: Final = APIRouter(prefix=MANAGEMENT_V1_PREFIX) + + +@router.post( + "/users/bulk", + tags=["Internal User management"], # mutable-ok: fastapi types tags as list[str | Enum] + dependencies=(Depends(user_api_key_auth),), + response_model=BulkNewUserResponse, +) +@management_endpoint_wrapper +async def bulk_create_users_route( + data: BulkNewUserRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> BulkNewUserResponse: + """ + Create up to 500 internal users in one request, optionally adding each one to teams. + + Every entry in `users` takes the same fields as `/user/new`, with two differences: `auto_create_key` + defaults to `false` (opt in per user to also get a virtual key back) and `send_invite_email` is not + supported. Unknown fields are rejected with 422. Rows are validated together (duplicate ids or emails, + unknown teams, roles the caller may not grant), inserted in one statement, and each referenced team is + written once for all of its new members. + + Rows fail independently: a bad row is reported in `data` with `success: false` and an `error`, and the + other rows still get created. A user that was created but could not be added to one of its teams is + reported with `success: true`, `teams` listing where they did land, and `error` naming the failed team. + The whole request is refused with a 403 problem document only if creating the valid rows would exceed + the license seat limit. + + Example curl: + ``` + curl -X POST "http://localhost:4000/management/v1/users/bulk" \\ + -H "Content-Type: application/json" \\ + -H "Authorization: Bearer sk-1234" \\ + -d '{ + "users": [ + {"user_email": "a@example.com", "user_role": "internal_user", "teams": ["team-1"]}, + {"user_email": "b@example.com", "user_role": "internal_user", "auto_create_key": true} + ] + }' + ``` + + Returns `data` (one entry per input row, in order, with `user_id`, `user_email`, `success`, `teams`, + `key`, `error`) and `meta` with `total_requested`, `created` and `failed`. + """ + try: + from litellm.proxy.proxy_server import ( + _license_check, # pyright: ignore[reportPrivateUsage] # same proxy license singleton /user/new reads + litellm_proxy_admin_name, + prisma_client, + user_api_key_cache, + ) + + if prisma_client is None: + raise ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}database-not-connected", + title="Database not connected", + status=503, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + ) + + return await bulk_create_users( + users=data.users, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + license_check=_license_check, + litellm_proxy_admin_name=litellm_proxy_admin_name, + user_api_key_cache=user_api_key_cache, + ) + + except ManagementProblem: + raise + except Exception: # noqa: BLE001 # a driver error answers as a problem document, not the OpenAI error shape + verbose_proxy_logger.exception("/management/v1/users/bulk: Exception occurred") + raise ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}internal-server-error", + title="Internal server error", + status=500, + detail="Failed to create users.", + ) + ) diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index 2eed5121755..56abe3b6a3f 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -1,4 +1,4 @@ -"""Batched internal user creation behind `/user/bulk_new`. +"""Batched internal user creation behind `POST /management/v1/users/bulk`. The batch is validated with set queries, user rows land in one `create_many`, and every referenced team is written once under its advisory lock instead of once per user. @@ -33,6 +33,7 @@ from litellm.proxy.auth.litellm_license import LicenseCheck from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks +from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE, ManagementProblem from litellm.proxy.management_endpoints.common_utils import ( _is_user_org_admin_for_team, # pyright: ignore[reportPrivateUsage] # same team-admin check /user/new uses _is_user_team_admin, # pyright: ignore[reportPrivateUsage] # same team-admin check /user/new uses @@ -61,9 +62,11 @@ from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkNewUserItem, + BulkNewUserMeta, BulkNewUserResponse, UserCreateResult, ) +from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail if TYPE_CHECKING: from prisma import Prisma @@ -776,7 +779,8 @@ async def bulk_create_users( ) -> BulkNewUserResponse: """Create every valid row in `users`; rows that fail validation or a write are reported, not raised. - Raises `HTTPException(403)` only when the whole batch would push the deployment over its license seat limit. + Raises a 403 `ManagementProblem` only when the whole batch would push the deployment over its license seat + limit. """ pending, request_failures = _partition_rows(users, user_api_key_dict) existing_ids, existing_emails = await _existing_user_conflicts(prisma_client, pending) @@ -791,9 +795,13 @@ async def bulk_create_users( billable_users: Final = await UserRepository(prisma_client).count_billable_users() if creatable and license_check.is_over_limit(total_users=billable_users + len(creatable)): - raise HTTPException( - status_code=403, - detail="License is over limit. Please contact support@berri.ai to upgrade your license.", + raise ManagementProblem( + ProblemDetail( + type=f"{PROBLEM_TYPE_BASE}license-limit-exceeded", + title="License limit exceeded", + status=403, + detail="License is over limit. Please contact support@berri.ai to upgrade your license.", + ) ) prepared_outcomes: Final = tuple([await _prepare_user(user, prisma_client) for user in creatable]) @@ -858,8 +866,6 @@ async def bulk_create_users( ) successes: Final = sum(1 for result in results if result.success) return BulkNewUserResponse( - results=results, - total_requested=len(users), - successful_creations=successes, - failed_creations=len(users) - successes, + data=results, + meta=BulkNewUserMeta(total_requested=len(users), created=successes, failed=len(users) - successes), ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d9f8e04ebda..367f4b02a4c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -473,9 +473,10 @@ from litellm.proxy.hooks.prompt_injection_detection import ( from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger, run_spend_event from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.list_api.common import ( - PROBLEM_TYPE_BASE, ManagementProblem, + ValidationErrorDetail, problem_response, + request_validation_problem, ) from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( @@ -598,7 +599,6 @@ from litellm.proxy.spend_tracking.spend_event_producer import ( SpendEventProducer, build_spend_event_producer, ) -from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail try: from litellm.proxy.enterprise_billing.billing_metrics import ( @@ -1784,27 +1784,13 @@ class _ExceptionRow(TypedDict, total=False): exception_counts: Mapping[str, int] -class _ValidationErrorDetail(TypedDict): - loc: tuple[int | str, ...] - msg: str - - @app.exception_handler(RequestValidationError) async def otel_request_validation_exception_handler(request: Request, exc: RequestValidationError): if request.url.path.startswith(MANAGEMENT_V1_PREFIX): - _close_dangling_otel_server_span(request, 400, exc=exc) - validation_errors: Final[Sequence[_ValidationErrorDetail]] = exc.errors() - return problem_response( - ProblemDetail( - type=f"{PROBLEM_TYPE_BASE}invalid-query-parameter", - title="Invalid query parameter", - status=400, - detail="; ".join( - f"{'.'.join(str(part) for part in error['loc'][1:])}: {error['msg']}" for error in validation_errors - ) - or "The request query parameters are invalid.", - ) - ) + validation_errors: Final[Sequence[ValidationErrorDetail]] = exc.errors() + problem: Final = request_validation_problem(validation_errors) + _close_dangling_otel_server_span(request, problem.status, exc=exc) + return problem_response(problem) _close_dangling_otel_server_span(request, 422, exc=exc) return JSONResponse( status_code=422, diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index a7e5c6d2916..f6a2c173ad5 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -1,7 +1,7 @@ from collections.abc import Mapping from typing import Any, Final, Literal -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator from typing_extensions import ReadOnly, TypedDict from litellm.proxy._types import ( @@ -89,7 +89,10 @@ class BulkUpdateUserResponse(BaseModel): class BulkNewUserItem(NewUserRequest): - """One row of `/user/bulk_new`: the `/user/new` body, with keys opt-in and invite emails unsupported.""" + """One row of `POST /management/v1/users/bulk`: the `/user/new` body, with keys opt-in and invite emails + unsupported. Unknown fields are rejected, as on every `/management/v1` request body.""" + + model_config = ConfigDict(extra="forbid", protected_namespaces=()) auto_create_key: bool = False @@ -97,16 +100,19 @@ class BulkNewUserItem(NewUserRequest): @classmethod def reject_invite_email(cls, value: bool | None) -> bool | None: if value: - raise ValueError("send_invite_email is not supported on /user/bulk_new; invite users separately") + raise ValueError("send_invite_email is not supported on /management/v1/users/bulk; invite users separately") return value class BulkNewUserRequest(BaseModel): + model_config = ConfigDict(extra="forbid") + users: tuple[BulkNewUserItem, ...] = Field(min_length=1, max_length=MAX_BULK_NEW_USERS) class UserCreateResult(BaseModel): - """Outcome for one row of `/user/bulk_new`. `teams` lists the teams the user was actually added to.""" + """Outcome for one row of `POST /management/v1/users/bulk`. `teams` lists the teams the user was actually + added to.""" user_id: str | None = None user_email: str | None = None @@ -116,8 +122,14 @@ class UserCreateResult(BaseModel): error: str | None = None -class BulkNewUserResponse(BaseModel): - results: tuple[UserCreateResult, ...] +class BulkNewUserMeta(BaseModel): total_requested: int - successful_creations: int - failed_creations: int + created: int + failed: int + + +class BulkNewUserResponse(BaseModel): + """`data` holds one result per input row, in input order.""" + + data: tuple[UserCreateResult, ...] + meta: BulkNewUserMeta diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 6f85df98850..6bc8947e89f 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -84,7 +84,6 @@ POST /team/{team_id}/member/{user_id}/reset_spend POST /team/key/bulk_update POST /team/permissions_bulk_update POST /team/{team_id}/disable_logging -POST /user/bulk_new POST /user/bulk_update # Alternate method or path for functionality the provider already manages elsewhere diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py new file mode 100644 index 00000000000..b61e639e453 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py @@ -0,0 +1,122 @@ +"""The HTTP contract of `POST /management/v1/users/bulk`: envelope, problem documents and strict bodies. + +The batching behaviour itself is covered next to the helper, in +`tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py`, whose in-memory Prisma this reuses. +""" + +import pytest +from fastapi import FastAPI, Request +from fastapi.exceptions import RequestValidationError +from fastapi.testclient import TestClient + +from litellm.proxy._types import LitellmUserRoles, Member +from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth +from litellm.proxy.list_api.common import ManagementProblem, problem_response, request_validation_problem +from litellm.proxy.management_endpoints.management_v1 import router +from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX +from tests.test_litellm.proxy.management_helpers.test_bulk_user_creation import _FakePrisma, _License, _team + +app = FastAPI() + + +@app.exception_handler(ManagementProblem) +async def management_problem_exception_handler(request: Request, exc: ManagementProblem): + return problem_response(exc.problem) + + +@app.exception_handler(RequestValidationError) +async def validation_exception_handler(request: Request, exc: RequestValidationError): + return problem_response(request_validation_problem(exc.errors())) + + +app.include_router(router) +client = TestClient(app) + +USERS_BULK_PATH = f"{MANAGEMENT_V1_PREFIX}/users/bulk" + + +@pytest.fixture +def as_proxy_admin(): + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + yield + app.dependency_overrides.clear() + + +@pytest.fixture +def prisma(monkeypatch): + fake = _FakePrisma(teams=[_team("t1", [Member(user_id="existing", role="admin")])]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check", _License()) + return fake + + +def _post(body: object): + return client.post(USERS_BULK_PATH, json=body, headers={"Authorization": "Bearer k"}) + + +def test_returns_one_result_per_row_in_order_inside_the_data_meta_envelope(prisma, as_proxy_admin): + response = _post( + { + "users": [ + {"user_id": "u1", "user_email": "a@example.com", "teams": ["t1"]}, + {"user_id": "u2", "teams": ["missing-team"]}, + {"user_id": "u3"}, + ] + } + ) + + assert response.status_code == 200 + body = response.json() + assert set(body) == {"data", "meta"} + assert body["meta"] == {"total_requested": 3, "created": 2, "failed": 1} + assert [row["user_id"] for row in body["data"]] == ["u1", "u2", "u3"] + assert [row["success"] for row in body["data"]] == [True, False, True] + assert body["data"][0]["teams"] == ["t1"] + assert "missing-team" in body["data"][1]["error"] + assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["existing", "u1"] + + +def test_an_unknown_field_anywhere_in_the_body_is_a_422_problem(prisma, as_proxy_admin): + for body in ( + {"users": [{"user_email": "a@example.com", "user_emial": "typo"}]}, + {"users": [{"user_email": "a@example.com"}], "dry_run": True}, + ): + response = _post(body) + + assert response.status_code == 422, body + assert response.headers["content-type"] == "application/problem+json" + assert response.json()["type"] == "urn:litellm:error:invalid-request-body" + assert "Extra inputs are not permitted" in response.json()["detail"] + assert prisma.db.litellm_usertable.rows == {} + + +def test_empty_and_oversized_batches_are_422_problems(prisma, as_proxy_admin): + for users in ([], [{"user_email": f"{i}@example.com"} for i in range(501)]): + response = _post({"users": users}) + + assert response.status_code == 422, len(users) + assert response.json()["type"] == "urn:litellm:error:invalid-request-body" + assert prisma.db.litellm_usertable.rows == {} + + +def test_license_limit_is_a_403_problem_and_creates_nothing(prisma, as_proxy_admin, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server._license_check", _License(max_users=1)) + + response = _post({"users": [{"user_id": "u1"}, {"user_id": "u2"}]}) + + assert response.status_code == 403 + assert response.headers["content-type"] == "application/problem+json" + assert response.json()["type"] == "urn:litellm:error:license-limit-exceeded" + assert prisma.db.litellm_usertable.rows == {} + + +def test_no_database_is_a_503_problem(as_proxy_admin, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + + response = _post({"users": [{"user_id": "u1"}]}) + + assert response.status_code == 503 + assert response.headers["content-type"] == "application/problem+json" + assert response.json()["type"] == "urn:litellm:error:database-not-connected" diff --git a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py index 16e56c1b723..b5349fc2387 100644 --- a/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py @@ -4,12 +4,12 @@ from typing import Final import httpx import pytest -from fastapi import HTTPException from prisma.errors import UniqueViolationError from pydantic import BaseModel, ConfigDict, ValidationError from litellm.caching.caching import DualCache from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, UserAPIKeyAuth +from litellm.proxy.list_api.common import ManagementProblem from litellm.proxy.management_helpers.bulk_user_creation import bulk_create_users from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkNewUserItem, @@ -203,10 +203,10 @@ async def test_creates_users_and_team_membership_in_every_store(): ], ) - assert (response.total_requested, response.successful_creations, response.failed_creations) == (3, 3, 0) - assert [r.user_id for r in response.results] == ["u1", "u2", "u3"] - assert all(r.success and r.key is None and r.error is None for r in response.results) - assert [r.teams for r in response.results] == [("t1", "t2"), ("t1",), ()] + assert (response.meta.total_requested, response.meta.created, response.meta.failed) == (3, 3, 0) + assert [r.user_id for r in response.data] == ["u1", "u2", "u3"] + assert all(r.success and r.key is None and r.error is None for r in response.data) + assert [r.teams for r in response.data] == [("t1", "t2"), ("t1",), ()] users = prisma.db.litellm_usertable.rows assert users["u1"].teams == ["t1", "t2"] and users["u1"].max_budget == 50 @@ -225,9 +225,9 @@ async def test_user_id_already_on_the_roster_keeps_the_team_and_is_not_added_twi prisma = _FakePrisma(teams=[_team("t1", [Member(user_id="u1", role="user")])]) response = await _run(prisma, [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2", "teams": ["t1"]}]) - assert [r.success for r in response.results] == [True, True] - assert [r.teams for r in response.results] == [("t1",), ("t1",)] - assert [r.error for r in response.results] == [None, None] + assert [r.success for r in response.data] == [True, True] + assert [r.teams for r in response.data] == [("t1",), ("t1",)] + assert [r.error for r in response.data] == [None, None] assert prisma.db.litellm_usertable.rows["u1"].teams == ["t1"] assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1", "u2"] @@ -267,9 +267,9 @@ async def test_bad_rows_fail_alone_and_good_rows_still_land(): ], ) - assert [r.success for r in response.results] == [True, False, False, False, False, False, False, False, True] - assert (response.successful_creations, response.failed_creations) == (2, 7) - errors = [r.error for r in response.results] + assert [r.success for r in response.data] == [True, False, False, False, False, False, False, False, True] + assert (response.meta.created, response.meta.failed) == (2, 7) + errors = [r.error for r in response.data] assert "Duplicate user_email" in errors[1] assert "Duplicate user_id" in errors[2] assert "already exists" in errors[3] and "already exists" in errors[4] @@ -289,8 +289,8 @@ async def test_insert_failure_falls_back_to_per_row_and_reports_only_that_row(): [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2", "teams": ["t1"]}, {"user_id": "u3"}], ) - assert [r.success for r in response.results] == [True, False, True] - assert "insert failed for u2" in (response.results[1].error or "") + assert [r.success for r in response.data] == [True, False, True] + assert "insert failed for u2" in (response.data[1].error or "") assert set(prisma.db.litellm_usertable.rows) == {"u1", "u3"} assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1"] @@ -300,8 +300,8 @@ async def test_insert_that_committed_but_lost_its_response_still_counts_as_creat prisma = _FakePrisma(teams=[_team("t1")], commit_then_drop=True) response = await _run(prisma, [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2"}]) - assert [r.success for r in response.results] == [True, True] - assert [r.error for r in response.results] == [None, None] + assert [r.success for r in response.data] == [True, True] + assert [r.error for r in response.data] == [None, None] assert set(prisma.db.litellm_usertable.rows) == {"u1", "u2"} assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1"] @@ -311,8 +311,8 @@ async def test_user_id_taken_by_a_concurrent_request_is_not_claimed_by_this_batc prisma = _FakePrisma(teams=[_team("t1")], raced_ids=frozenset({"u1"})) response = await _run(prisma, [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2", "teams": ["t1"]}]) - assert [r.success for r in response.results] == [False, True] - assert "User id=u1 already exists" in (response.results[0].error or "") + assert [r.success for r in response.data] == [False, True] + assert "User id=u1 already exists" in (response.data[0].error or "") assert prisma.db.litellm_usertable.rows["u1"].user_email == "u1@other-request.example" assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u2"] @@ -327,12 +327,12 @@ async def test_team_write_failure_keeps_user_and_reports_it_on_the_row(): prisma.db.litellm_teamtable.update = explode response = await _run(prisma, [{"user_id": "u1", "teams": ["t1", "t2"]}]) - result = response.results[0] + result = response.data[0] assert result.success is True assert result.teams == () assert "t1" in (result.error or "") and "roster write failed" in (result.error or "") assert prisma.db.litellm_usertable.rows["u1"].teams == [] - assert (response.successful_creations, response.failed_creations) == (1, 0) + assert (response.meta.created, response.meta.failed) == (1, 0) @pytest.mark.asyncio @@ -364,7 +364,7 @@ async def test_keys_are_opt_in_per_row(): generate_key=generate_key, ) - assert [r.key for r in response.results] == [None, "sk-u2", None] + assert [r.key for r in response.data] == [None, "sk-u2", None] assert len(calls) == 1 assert calls[0]["user_id"] == "u2" and calls[0]["table_name"] == "key" assert calls[0]["models"] == ("gpt-4o",) and calls[0]["key_alias"] == "u2-key" @@ -385,8 +385,8 @@ async def test_non_admin_cannot_create_admin_users_but_other_rows_proceed(): caller=INTERNAL, ) - assert [r.success for r in response.results] == [False, True] - assert "Only proxy admins" in (response.results[0].error or "") + assert [r.success for r in response.data] == [False, True] + assert "Only proxy admins" in (response.data[0].error or "") assert set(prisma.db.litellm_usertable.rows) == {"u2"} @@ -396,19 +396,19 @@ async def test_license_is_checked_once_against_the_whole_batch(): prisma.db.litellm_usertable.rows["existing"] = _UserRow(user_id="existing") license = _License(max_users=3) - with pytest.raises(HTTPException) as exc: + with pytest.raises(ManagementProblem) as exc: await _run(prisma, [{"user_id": f"u{i}"} for i in range(3)], license=license) - assert exc.value.status_code == 403 + assert (exc.value.problem.status, exc.value.problem.type) == (403, "urn:litellm:error:license-limit-exceeded") assert license.seen == [4] assert set(prisma.db.litellm_usertable.rows) == {"existing"} ok = await _run(prisma, [{"user_id": f"u{i}"} for i in range(2)], license=license) - assert ok.successful_creations == 2 + assert ok.meta.created == 2 resend = await _run(prisma, [{"user_id": f"u{i}"} for i in range(2)], license=license) - assert [r.success for r in resend.results] == [False, False] - assert all("already exists" in (r.error or "") for r in resend.results) + assert [r.success for r in resend.data] == [False, False] + assert all("already exists" in (r.error or "") for r in resend.data) assert license.seen == [4, 3] assert set(prisma.db.litellm_usertable.rows) == {"existing", "u0", "u1"} @@ -422,3 +422,10 @@ def test_request_rejects_empty_oversized_and_invite_rows(): BulkNewUserItem(user_email="a@example.com", send_invite_email=True) assert len(BulkNewUserRequest(users=[{"user_email": f"{i}@example.com"} for i in range(500)]).users) == 500 assert BulkNewUserItem(user_email="a@example.com").auto_create_key is False + + +def test_request_rejects_unknown_fields_at_both_levels(): + with pytest.raises(ValidationError, match="extra_forbidden"): + BulkNewUserRequest(users=[{"user_email": "a@example.com", "user_emial": "typo"}]) + with pytest.raises(ValidationError, match="extra_forbidden"): + BulkNewUserRequest(users=[{"user_email": "a@example.com"}], dry_run=True) diff --git a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py index 4aea2e16364..53ea761daa7 100644 --- a/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py +++ b/tests/test_litellm/proxy/proxy_server/test_exception_handlers.py @@ -227,7 +227,9 @@ async def test_otel_request_validation_exception_handler_empty_errors_invalid_pa async def test_otel_request_validation_exception_handler_returns_a_problem_on_the_control_plane(): """`/management/v1` answers validation errors as RFC 9457, so a caller there gets a 400 problem document rather than the proxy-wide 422 `{"detail": [...]}` shape.""" - errors = [{"loc": ["query", "page_size"], "msg": "Input should be less than or equal to 100", "type": "less_than_equal"}] + errors = [ + {"loc": ["query", "page_size"], "msg": "Input should be less than or equal to 100", "type": "less_than_equal"} + ] exc = RequestValidationError(errors) request = _make_request(path="/management/v1/spend_logs/end_users") @@ -242,6 +244,26 @@ async def test_otel_request_validation_exception_handler_returns_a_problem_on_th assert "detail" in body and not isinstance(body["detail"], list) +@pytest.mark.asyncio +async def test_otel_request_validation_exception_handler_answers_a_bad_control_plane_body_with_422(): + """A request body that fails validation, an unknown field included, is 422 on + `/management/v1`; only query parameter problems are 400.""" + errors = [ + {"loc": ["body", "users", 0, "user_emial"], "msg": "Extra inputs are not permitted", "type": "extra_forbidden"} + ] + exc = RequestValidationError(errors) + request = _make_request(path="/management/v1/users/bulk") + + response = await otel_request_validation_exception_handler(request=request, exc=exc) + body = json.loads(response.body) + + assert response.status_code == 422 + assert response.media_type == "application/problem+json" + assert body["type"] == "urn:litellm:error:invalid-request-body" + assert body["status"] == 422 + assert "users.0.user_emial: Extra inputs are not permitted" in body["detail"] + + @pytest.mark.asyncio async def test_otel_request_validation_exception_handler_leaves_other_routes_on_422(): """The problem+json branch is scoped by path prefix. A route that merely contains @@ -249,9 +271,7 @@ async def test_otel_request_validation_exception_handler_leaves_other_routes_on_ exc = RequestValidationError([]) for path in ("/management", "/v1/management/foo", "/customer/list"): - response = await otel_request_validation_exception_handler( - request=_make_request(path=path), exc=exc - ) + response = await otel_request_validation_exception_handler(request=_make_request(path=path), exc=exc) assert response.status_code == 422, path assert json.loads(response.body) == {"detail": []}, path @@ -294,6 +314,4 @@ async def test_otel_unhandled_exception_handler_reraises_proxy_exception_error() async def test_otel_unhandled_exception_handler_reraises_http_exception_invalid(): request = _make_request() with pytest.raises(HTTPException): - await otel_unhandled_exception_handler( - request=request, exc=HTTPException(status_code=418, detail="teapot") - ) + await otel_unhandled_exception_handler(request=request, exc=HTTPException(status_code=418, detail="teapot")) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f63bed3c3ec..366edbe13e0 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8470,6 +8470,54 @@ export interface paths { patch?: never; trace?: never; }; + "/management/v1/users/bulk": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Bulk Create Users Route + * @description Create up to 500 internal users in one request, optionally adding each one to teams. + * + * Every entry in `users` takes the same fields as `/user/new`, with two differences: `auto_create_key` + * defaults to `false` (opt in per user to also get a virtual key back) and `send_invite_email` is not + * supported. Unknown fields are rejected with 422. Rows are validated together (duplicate ids or emails, + * unknown teams, roles the caller may not grant), inserted in one statement, and each referenced team is + * written once for all of its new members. + * + * Rows fail independently: a bad row is reported in `data` with `success: false` and an `error`, and the + * other rows still get created. A user that was created but could not be added to one of its teams is + * reported with `success: true`, `teams` listing where they did land, and `error` naming the failed team. + * The whole request is refused with a 403 problem document only if creating the valid rows would exceed + * the license seat limit. + * + * Example curl: + * ``` + * curl -X POST "http://localhost:4000/management/v1/users/bulk" \ + * -H "Content-Type: application/json" \ + * -H "Authorization: Bearer sk-1234" \ + * -d '{ + * "users": [ + * {"user_email": "a@example.com", "user_role": "internal_user", "teams": ["team-1"]}, + * {"user_email": "b@example.com", "user_role": "internal_user", "auto_create_key": true} + * ] + * }' + * ``` + * + * Returns `data` (one entry per input row, in order, with `user_id`, `user_email`, `success`, `teams`, + * `key`, `error`) and `meta` with `total_requested`, `created` and `failed`. + */ + post: operations["bulk_create_users_route_management_v1_users_bulk_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/mcp": { parameters: { query?: never; @@ -16478,53 +16526,6 @@ export interface paths { patch?: never; trace?: never; }; - "/user/bulk_new": { - parameters: { - query?: never; - header?: never; - path?: never; - cookie?: never; - }; - get?: never; - put?: never; - /** - * Bulk New User - * @description Create up to 500 internal users in one request, optionally adding each one to teams. - * - * Every entry in `users` takes the same fields as `/user/new`, with two differences: `auto_create_key` - * defaults to `false` (opt in per user to also get a virtual key back) and `send_invite_email` is not - * supported. Rows are validated together (duplicate ids or emails, unknown teams, roles the caller may not - * grant), inserted in one statement, and each referenced team is written once for all of its new members. - * - * Rows fail independently: a bad row is reported in `results` with `success: false` and an `error`, and the - * other rows still get created. A user that was created but could not be added to one of its teams is - * reported with `success: true`, `teams` listing where they did land, and `error` naming the failed team. - * The whole request is rejected with 403 only if creating the valid rows would exceed the license seat limit. - * - * Usage Example - * - * ```shell - * curl -X POST "http://localhost:4000/user/bulk_new" \ - * -H "Content-Type: application/json" \ - * -H "Authorization: Bearer sk-1234" \ - * -d '{ - * "users": [ - * {"user_email": "a@example.com", "user_role": "internal_user", "teams": ["team-1"]}, - * {"user_email": "b@example.com", "user_role": "internal_user", "auto_create_key": true} - * ] - * }' - * ``` - * - * Returns `results` (one entry per input row, in order, with `user_id`, `user_email`, `success`, `teams`, - * `key`, `error`), `total_requested`, `successful_creations` and `failed_creations`. - */ - post: operations["bulk_new_user_user_bulk_new_post"]; - delete?: never; - options?: never; - head?: never; - patch?: never; - trace?: never; - }; "/user/bulk_update": { parameters: { query?: never; @@ -24552,7 +24553,8 @@ export interface components { }; /** * BulkNewUserItem - * @description One row of `/user/bulk_new`: the `/user/new` body, with keys opt-in and invite emails unsupported. + * @description One row of `POST /management/v1/users/bulk`: the `/user/new` body, with keys opt-in and invite emails + * unsupported. Unknown fields are rejected, as on every `/management/v1` request body. */ BulkNewUserItem: { /** Agent Id */ @@ -24676,21 +24678,28 @@ export interface components { /** User Role */ user_role?: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; + /** BulkNewUserMeta */ + BulkNewUserMeta: { + /** Created */ + created: number; + /** Failed */ + failed: number; + /** Total Requested */ + total_requested: number; + }; /** BulkNewUserRequest */ BulkNewUserRequest: { /** Users */ users: components["schemas"]["BulkNewUserItem"][]; }; - /** BulkNewUserResponse */ + /** + * BulkNewUserResponse + * @description `data` holds one result per input row, in input order. + */ BulkNewUserResponse: { - /** Failed Creations */ - failed_creations: number; - /** Results */ - results: components["schemas"]["UserCreateResult"][]; - /** Successful Creations */ - successful_creations: number; - /** Total Requested */ - total_requested: number; + /** Data */ + data: components["schemas"]["UserCreateResult"][]; + meta: components["schemas"]["BulkNewUserMeta"]; }; /** * BulkTeamMemberAddRequest @@ -39534,7 +39543,8 @@ export interface components { }; /** * UserCreateResult - * @description Outcome for one row of `/user/bulk_new`. `teams` lists the teams the user was actually added to. + * @description Outcome for one row of `POST /management/v1/users/bulk`. `teams` lists the teams the user was actually + * added to. */ UserCreateResult: { /** Error */ @@ -51322,6 +51332,39 @@ export interface operations { }; }; }; + bulk_create_users_route_management_v1_users_bulk_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["BulkNewUserRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["BulkNewUserResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; aggregate_mcp_route_mcp_get: { parameters: { query?: never; @@ -60836,39 +60879,6 @@ export interface operations { }; }; }; - bulk_new_user_user_bulk_new_post: { - parameters: { - query?: never; - header?: never; - path?: never; - cookie?: never; - }; - requestBody: { - content: { - "application/json": components["schemas"]["BulkNewUserRequest"]; - }; - }; - responses: { - /** @description Successful Response */ - 200: { - headers: { - [name: string]: unknown; - }; - content: { - "application/json": components["schemas"]["BulkNewUserResponse"]; - }; - }; - /** @description Validation Error */ - 422: { - headers: { - [name: string]: unknown; - }; - content: { - "application/json": components["schemas"]["HTTPValidationError"]; - }; - }; - }; - }; bulk_user_update_user_bulk_update_post: { parameters: { query?: never; From f78dd921c958c38c3adba9ee490071a658c79e27 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 22:34:33 +0000 Subject: [PATCH 059/107] fix(guardrails): keep legacy not_run neutral and stop labelling image-only input as skipped Usage tracking, compliance and the dashboard now treat both not_run (older spend logs) and skipped as unevaluated through a shared UNEVALUATED_GUARDRAIL_STATUSES set, so old records stop counting as passed. The skipped record is no longer written when the request carried images, since images without text were never dispatched to guardrails before this change and that gap is not a message-scoping skip Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/guardrail_translation/handler.py | 2 +- litellm/proxy/compliance_checks.py | 7 ++++- litellm/proxy/guardrails/usage_tracking.py | 3 +- litellm/types/utils.py | 2 ++ .../test_openai_guardrail_handler.py | 21 +++++++++++++ .../proxy/guardrails/test_usage_tracking.py | 12 ++++--- .../test_compliance_endpoints.py | 5 +-- .../GuardrailViewer/GuardrailViewer.test.tsx | 31 ++++++++++--------- .../GuardrailViewer/GuardrailViewer.tsx | 2 +- .../LogDetailsDrawer/LogDetailContent.tsx | 2 +- 10 files changed, 61 insertions(+), 26 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index dca90e07421..035dce46d27 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -210,7 +210,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings=tool_call_task_mappings, ) - elif not guardrail_to_apply.records_own_guardrail_information: + elif not images_to_check and not guardrail_to_apply.records_own_guardrail_information: guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response="no scannable content after message scoping", request_data=data, diff --git a/litellm/proxy/compliance_checks.py b/litellm/proxy/compliance_checks.py index ef2d8fb6e20..18123156411 100644 --- a/litellm/proxy/compliance_checks.py +++ b/litellm/proxy/compliance_checks.py @@ -11,6 +11,7 @@ from litellm.types.proxy.compliance_endpoints import ( ComplianceCheckRequest, ComplianceCheckResult, ) +from litellm.types.utils import UNEVALUATED_GUARDRAIL_STATUSES class ComplianceChecker: @@ -26,7 +27,11 @@ class ComplianceChecker: def __init__(self, data: ComplianceCheckRequest): self.data = data - self.guardrails = [g for g in (data.guardrail_information or []) if g.get("guardrail_status") != "skipped"] + self.guardrails = [ + g + for g in (data.guardrail_information or []) + if g.get("guardrail_status") not in UNEVALUATED_GUARDRAIL_STATUSES + ] def _get_guardrails_by_mode(self, mode: str) -> list[dict]: """ diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 8a131fbfcde..789b8febfbf 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -27,6 +27,7 @@ from litellm.repositories.table_repositories import ( DailyGuardrailUsageUnitsRepository, SpendLogGuardrailIndexRepository, ) +from litellm.types.utils import UNEVALUATED_GUARDRAIL_STATUSES if TYPE_CHECKING: from prisma import types as prisma_types @@ -197,7 +198,7 @@ def guardrail_status_to_action(status: str | None) -> str: if not status: return "passed" s: Final = (status or "").lower() - if s == "skipped": + if s in UNEVALUATED_GUARDRAIL_STATUSES: return "skipped" if "intervened" in s or "block" in s: return "blocked" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cf302b3f27f..6da44007fb4 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3123,6 +3123,8 @@ GuardrailStatus = Literal[ "success", "guardrail_flagged", "guardrail_intervened", "guardrail_failed_to_respond", "not_run", "skipped" ] +UNEVALUATED_GUARDRAIL_STATUSES: Final[frozenset[GuardrailStatus]] = frozenset({"not_run", "skipped"}) + # Fields on a guardrail record whose values can quote the caller's prompt: the payload sent to the # guardrail, the provider response that echoes it back, and the two first-party hooks that inline # prompt substrings (``block_code_execution`` and ``litellm_content_filter``). Every other field diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index c6d01db45de..7bfd50fa6f8 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1942,6 +1942,27 @@ class TestNoScannableContentRecordsSkipped: assert guardrail.last_inputs is not None assert all(e.get("guardrail_status") != "skipped" for e in self._recorded_entries(data)) + @pytest.mark.asyncio + async def test_image_only_content_is_not_reported_as_skipped(self): + """Images are only scanned alongside text, so an image-only request is a + pre-existing scan gap, not a message-scoping skip, and must not be labelled one""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="image-guardrail") + guardrail.skip_system_message_in_guardrail = True + data = { + "messages": [ + {"role": "system", "content": "SYSTEM-PROMPT"}, + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}], + }, + ] + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert self._recorded_entries(data) == [] + class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks""" diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index cb85eb8b5ec..ae883d3eaeb 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -350,15 +350,17 @@ async def test_zero_and_non_int_usage_counters_are_skipped(): @pytest.mark.asyncio -async def test_skipped_entries_are_indexed_but_not_counted_as_evaluations(): +@pytest.mark.parametrize("status", ["skipped", "not_run"]) +async def test_skipped_entries_are_indexed_but_not_counted_as_evaluations(status: str): """ LIT-6314 records a skipped entry when message scoping leaves a guardrail - nothing to scan. The guardrail never evaluated the request, so counting it - as a passed evaluation would inflate daily pass rates; it still gets an - index row so per-request drill-down finds the spend log. + nothing to scan (older spend logs spell it not_run). The guardrail never + evaluated the request, so counting it as a passed evaluation would inflate + daily pass rates; it still gets an index row so per-request drill-down + finds the spend log. """ prisma = _prisma() - logs = [_payload("r1", guardrail_status="skipped"), _payload("r2")] + logs = [_payload("r1", guardrail_status=status), _payload("r2")] await process_spend_logs_guardrail_usage(prisma, logs) diff --git a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py index 9587e3d95ee..bafe608bff8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py @@ -594,14 +594,15 @@ class TestModeMatching: class TestSkippedGuardrails: """LIT-6314 logs a skipped entry for a guardrail that message scoping left nothing to scan.""" - def test_skipped_alone_never_evidences_compliance(self): + @pytest.mark.parametrize("status", ["skipped", "not_run"]) + def test_skipped_alone_never_evidences_compliance(self, status: str): data = ComplianceCheckRequest( request_id="req-601", user_id="user-1", model="gpt-4", timestamp="2026-02-17T00:00:00Z", guardrail_information=[ - {"guardrail_name": "pii_detection", "guardrail_status": "skipped", "guardrail_mode": "pre_call"}, + {"guardrail_name": "pii_detection", "guardrail_status": status, "guardrail_mode": "pre_call"}, ], ) results = {c.check_name: c.passed for c in ComplianceChecker(data).check_eu_ai_act()} diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx index 740ef857690..02a4c8008a6 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx @@ -68,22 +68,25 @@ describe("GuardrailViewer", () => { expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); }); - it("renders skipped as SKIPPED (muted) and keeps it out of the evaluated and passed counts", async () => { - const user = userEvent.setup(); - const data = makeGuardrailInformation(skippedPreCall); - renderWithProviders(); + it.each(["skipped", "not_run"])( + "renders %s as SKIPPED (muted) and keeps it out of the evaluated and passed counts", + async (guardrail_status) => { + const user = userEvent.setup(); + const data = makeGuardrailInformation({ ...skippedPreCall, guardrail_status }); + renderWithProviders(); - expect(screen.getByText(/0 guardrails evaluated/)).toBeInTheDocument(); - expect(screen.getByText(/0 Passed/)).toHaveClass("text-muted-foreground"); - expect(screen.getByText(/1 Skipped/)).toBeInTheDocument(); - const badge = screen.getByText("SKIPPED"); - expect(badge).toHaveClass("text-muted-foreground"); - expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); - expect(screen.queryByText(/^T\+/)).not.toBeInTheDocument(); + expect(screen.getByText(/0 guardrails evaluated/)).toBeInTheDocument(); + expect(screen.getByText(/0 Passed/)).toHaveClass("text-muted-foreground"); + expect(screen.getByText(/1 Skipped/)).toBeInTheDocument(); + const badge = screen.getByText("SKIPPED"); + expect(badge).toHaveClass("text-muted-foreground"); + expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); + expect(screen.queryByText(/^T\+/)).not.toBeInTheDocument(); - await user.click(screen.getByText("pii-rail")); - expect(screen.getByText("no scannable content after message scoping")).toBeInTheDocument(); - }); + await user.click(screen.getByText("pii-rail")); + expect(screen.getByText("no scannable content after message scoping")).toBeInTheDocument(); + }, + ); it("anchors the lifecycle timeline on timed entries when an untimed skipped entry sorts first", () => { const skipped = makeGuardrailInformation({ ...skippedPreCall, guardrail_name: "skipped-rail" }); diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx index c67d682a233..d171b7ff4b3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx @@ -140,7 +140,7 @@ const getEntryOutcome = (entry: GuardrailInformation): EntryOutcome => { const status = (entry.guardrail_status ?? "").toLowerCase(); if (status === "success") return "passed"; if (status === "guardrail_flagged") return "flagged"; - if (status === "skipped") return "skipped"; + if (status === "skipped" || status === "not_run") return "skipped"; return "failed"; }; diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx index e052faa2228..e022d166fe0 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx @@ -705,7 +705,7 @@ const GUARDRAIL_JUMP_LINK_STYLE = { const isPassedStatus = (status: unknown) => status === "pass" || status === "passed" || status === "success"; const isFlaggedStatus = (status: unknown) => status === "flagged" || status === "guardrail_flagged"; -const isSkippedStatus = (status: unknown) => status === "skipped"; +const isSkippedStatus = (status: unknown) => status === "skipped" || status === "not_run"; const guardrailJumpLinkOutcome = (evaluated: unknown[]): keyof typeof GUARDRAIL_JUMP_LINK_STYLE => { if (evaluated.length === 0) return "skipped"; From 8e83e91275d481df8d9a5e869d1c7675b1abf123 Mon Sep 17 00:00:00 2001 From: ryan Date: Mon, 14 Sep 2026 22:40:57 +0000 Subject: [PATCH 060/107] fix(proxy): keep bulk user row errors free of the tuple length message Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/internal_user_endpoints.py | 4 ++-- .../management_endpoints/management_v1/test_users.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index f6a2c173ad5..3b4a2a417f2 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator @@ -107,7 +107,7 @@ class BulkNewUserItem(NewUserRequest): class BulkNewUserRequest(BaseModel): model_config = ConfigDict(extra="forbid") - users: tuple[BulkNewUserItem, ...] = Field(min_length=1, max_length=MAX_BULK_NEW_USERS) + users: Sequence[BulkNewUserItem] = Field(min_length=1, max_length=MAX_BULK_NEW_USERS) class UserCreateResult(BaseModel): diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py index b61e639e453..edd1d315093 100644 --- a/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py @@ -79,16 +79,16 @@ def test_returns_one_result_per_row_in_order_inside_the_data_meta_envelope(prism def test_an_unknown_field_anywhere_in_the_body_is_a_422_problem(prisma, as_proxy_admin): - for body in ( - {"users": [{"user_email": "a@example.com", "user_emial": "typo"}]}, - {"users": [{"user_email": "a@example.com"}], "dry_run": True}, + for body, field in ( + ({"users": [{"user_email": "a@example.com", "user_emial": "typo"}]}, "users.0.user_emial"), + ({"users": [{"user_email": "a@example.com"}], "dry_run": True}, "dry_run"), ): response = _post(body) assert response.status_code == 422, body assert response.headers["content-type"] == "application/problem+json" assert response.json()["type"] == "urn:litellm:error:invalid-request-body" - assert "Extra inputs are not permitted" in response.json()["detail"] + assert response.json()["detail"] == f"{field}: Extra inputs are not permitted" assert prisma.db.litellm_usertable.rows == {} From bd9a87ea7683bfc8d547a8996e665a6631e82af9 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 23:46:09 +0000 Subject: [PATCH 061/107] Revert "fix(guardrails): keep legacy not_run neutral and stop labelling image-only input as skipped" This reverts commit f78dd921c958c38c3adba9ee490071a658c79e27. --- .../chat/guardrail_translation/handler.py | 2 +- litellm/proxy/compliance_checks.py | 7 +---- litellm/proxy/guardrails/usage_tracking.py | 3 +- litellm/types/utils.py | 2 -- .../test_openai_guardrail_handler.py | 21 ------------- .../proxy/guardrails/test_usage_tracking.py | 12 +++---- .../test_compliance_endpoints.py | 5 ++- .../GuardrailViewer/GuardrailViewer.test.tsx | 31 +++++++++---------- .../GuardrailViewer/GuardrailViewer.tsx | 2 +- .../LogDetailsDrawer/LogDetailContent.tsx | 2 +- 10 files changed, 26 insertions(+), 61 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 035dce46d27..dca90e07421 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -210,7 +210,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings=tool_call_task_mappings, ) - elif not images_to_check and not guardrail_to_apply.records_own_guardrail_information: + elif not guardrail_to_apply.records_own_guardrail_information: guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response="no scannable content after message scoping", request_data=data, diff --git a/litellm/proxy/compliance_checks.py b/litellm/proxy/compliance_checks.py index 18123156411..ef2d8fb6e20 100644 --- a/litellm/proxy/compliance_checks.py +++ b/litellm/proxy/compliance_checks.py @@ -11,7 +11,6 @@ from litellm.types.proxy.compliance_endpoints import ( ComplianceCheckRequest, ComplianceCheckResult, ) -from litellm.types.utils import UNEVALUATED_GUARDRAIL_STATUSES class ComplianceChecker: @@ -27,11 +26,7 @@ class ComplianceChecker: def __init__(self, data: ComplianceCheckRequest): self.data = data - self.guardrails = [ - g - for g in (data.guardrail_information or []) - if g.get("guardrail_status") not in UNEVALUATED_GUARDRAIL_STATUSES - ] + self.guardrails = [g for g in (data.guardrail_information or []) if g.get("guardrail_status") != "skipped"] def _get_guardrails_by_mode(self, mode: str) -> list[dict]: """ diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 789b8febfbf..8a131fbfcde 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -27,7 +27,6 @@ from litellm.repositories.table_repositories import ( DailyGuardrailUsageUnitsRepository, SpendLogGuardrailIndexRepository, ) -from litellm.types.utils import UNEVALUATED_GUARDRAIL_STATUSES if TYPE_CHECKING: from prisma import types as prisma_types @@ -198,7 +197,7 @@ def guardrail_status_to_action(status: str | None) -> str: if not status: return "passed" s: Final = (status or "").lower() - if s in UNEVALUATED_GUARDRAIL_STATUSES: + if s == "skipped": return "skipped" if "intervened" in s or "block" in s: return "blocked" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6da44007fb4..cf302b3f27f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3123,8 +3123,6 @@ GuardrailStatus = Literal[ "success", "guardrail_flagged", "guardrail_intervened", "guardrail_failed_to_respond", "not_run", "skipped" ] -UNEVALUATED_GUARDRAIL_STATUSES: Final[frozenset[GuardrailStatus]] = frozenset({"not_run", "skipped"}) - # Fields on a guardrail record whose values can quote the caller's prompt: the payload sent to the # guardrail, the provider response that echoes it back, and the two first-party hooks that inline # prompt substrings (``block_code_execution`` and ``litellm_content_filter``). Every other field diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 7bfd50fa6f8..c6d01db45de 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1942,27 +1942,6 @@ class TestNoScannableContentRecordsSkipped: assert guardrail.last_inputs is not None assert all(e.get("guardrail_status") != "skipped" for e in self._recorded_entries(data)) - @pytest.mark.asyncio - async def test_image_only_content_is_not_reported_as_skipped(self): - """Images are only scanned alongside text, so an image-only request is a - pre-existing scan gap, not a message-scoping skip, and must not be labelled one""" - handler = OpenAIChatCompletionsHandler() - guardrail = MockGuardrail(guardrail_name="image-guardrail") - guardrail.skip_system_message_in_guardrail = True - data = { - "messages": [ - {"role": "system", "content": "SYSTEM-PROMPT"}, - { - "role": "user", - "content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}], - }, - ] - } - - await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - - assert self._recorded_entries(data) == [] - class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks""" diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index ae883d3eaeb..cb85eb8b5ec 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -350,17 +350,15 @@ async def test_zero_and_non_int_usage_counters_are_skipped(): @pytest.mark.asyncio -@pytest.mark.parametrize("status", ["skipped", "not_run"]) -async def test_skipped_entries_are_indexed_but_not_counted_as_evaluations(status: str): +async def test_skipped_entries_are_indexed_but_not_counted_as_evaluations(): """ LIT-6314 records a skipped entry when message scoping leaves a guardrail - nothing to scan (older spend logs spell it not_run). The guardrail never - evaluated the request, so counting it as a passed evaluation would inflate - daily pass rates; it still gets an index row so per-request drill-down - finds the spend log. + nothing to scan. The guardrail never evaluated the request, so counting it + as a passed evaluation would inflate daily pass rates; it still gets an + index row so per-request drill-down finds the spend log. """ prisma = _prisma() - logs = [_payload("r1", guardrail_status=status), _payload("r2")] + logs = [_payload("r1", guardrail_status="skipped"), _payload("r2")] await process_spend_logs_guardrail_usage(prisma, logs) diff --git a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py index bafe608bff8..9587e3d95ee 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py @@ -594,15 +594,14 @@ class TestModeMatching: class TestSkippedGuardrails: """LIT-6314 logs a skipped entry for a guardrail that message scoping left nothing to scan.""" - @pytest.mark.parametrize("status", ["skipped", "not_run"]) - def test_skipped_alone_never_evidences_compliance(self, status: str): + def test_skipped_alone_never_evidences_compliance(self): data = ComplianceCheckRequest( request_id="req-601", user_id="user-1", model="gpt-4", timestamp="2026-02-17T00:00:00Z", guardrail_information=[ - {"guardrail_name": "pii_detection", "guardrail_status": status, "guardrail_mode": "pre_call"}, + {"guardrail_name": "pii_detection", "guardrail_status": "skipped", "guardrail_mode": "pre_call"}, ], ) results = {c.check_name: c.passed for c in ComplianceChecker(data).check_eu_ai_act()} diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx index 02a4c8008a6..740ef857690 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx @@ -68,25 +68,22 @@ describe("GuardrailViewer", () => { expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); }); - it.each(["skipped", "not_run"])( - "renders %s as SKIPPED (muted) and keeps it out of the evaluated and passed counts", - async (guardrail_status) => { - const user = userEvent.setup(); - const data = makeGuardrailInformation({ ...skippedPreCall, guardrail_status }); - renderWithProviders(); + it("renders skipped as SKIPPED (muted) and keeps it out of the evaluated and passed counts", async () => { + const user = userEvent.setup(); + const data = makeGuardrailInformation(skippedPreCall); + renderWithProviders(); - expect(screen.getByText(/0 guardrails evaluated/)).toBeInTheDocument(); - expect(screen.getByText(/0 Passed/)).toHaveClass("text-muted-foreground"); - expect(screen.getByText(/1 Skipped/)).toBeInTheDocument(); - const badge = screen.getByText("SKIPPED"); - expect(badge).toHaveClass("text-muted-foreground"); - expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); - expect(screen.queryByText(/^T\+/)).not.toBeInTheDocument(); + expect(screen.getByText(/0 guardrails evaluated/)).toBeInTheDocument(); + expect(screen.getByText(/0 Passed/)).toHaveClass("text-muted-foreground"); + expect(screen.getByText(/1 Skipped/)).toBeInTheDocument(); + const badge = screen.getByText("SKIPPED"); + expect(badge).toHaveClass("text-muted-foreground"); + expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); + expect(screen.queryByText(/^T\+/)).not.toBeInTheDocument(); - await user.click(screen.getByText("pii-rail")); - expect(screen.getByText("no scannable content after message scoping")).toBeInTheDocument(); - }, - ); + await user.click(screen.getByText("pii-rail")); + expect(screen.getByText("no scannable content after message scoping")).toBeInTheDocument(); + }); it("anchors the lifecycle timeline on timed entries when an untimed skipped entry sorts first", () => { const skipped = makeGuardrailInformation({ ...skippedPreCall, guardrail_name: "skipped-rail" }); diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx index d171b7ff4b3..c67d682a233 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx @@ -140,7 +140,7 @@ const getEntryOutcome = (entry: GuardrailInformation): EntryOutcome => { const status = (entry.guardrail_status ?? "").toLowerCase(); if (status === "success") return "passed"; if (status === "guardrail_flagged") return "flagged"; - if (status === "skipped" || status === "not_run") return "skipped"; + if (status === "skipped") return "skipped"; return "failed"; }; diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx index e022d166fe0..e052faa2228 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx @@ -705,7 +705,7 @@ const GUARDRAIL_JUMP_LINK_STYLE = { const isPassedStatus = (status: unknown) => status === "pass" || status === "passed" || status === "success"; const isFlaggedStatus = (status: unknown) => status === "flagged" || status === "guardrail_flagged"; -const isSkippedStatus = (status: unknown) => status === "skipped" || status === "not_run"; +const isSkippedStatus = (status: unknown) => status === "skipped"; const guardrailJumpLinkOutcome = (evaluated: unknown[]): keyof typeof GUARDRAIL_JUMP_LINK_STYLE => { if (evaluated.length === 0) return "skipped"; From 0d0b96ed0670a505c361b2da42605dc4113c4dce Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 23:46:09 +0000 Subject: [PATCH 062/107] Revert "refactor(guardrails): rename scoped-out evaluation status from not_run to skipped" This reverts commit b37ce94075124b4429996cd52313d100fbb5212e. --- litellm/litellm_core_utils/litellm_logging.py | 1 - .../chat/guardrail_translation/handler.py | 2 +- litellm/proxy/compliance_checks.py | 2 +- litellm/proxy/guardrails/usage_endpoints.py | 4 ++-- litellm/proxy/guardrails/usage_tracking.py | 8 +++---- litellm/types/utils.py | 3 +-- .../test_litellm_logging.py | 12 ---------- .../test_openai_guardrail_handler.py | 8 +++---- .../proxy/guardrails/test_usage_endpoints.py | 10 ++++---- .../proxy/guardrails/test_usage_tracking.py | 22 ++++++++--------- .../test_compliance_endpoints.py | 12 +++++----- .../GuardrailsMonitor/LogViewer.test.tsx | 8 +++---- .../GuardrailsMonitor/LogViewer.tsx | 6 ++--- .../components/GuardrailsMonitor/mockData.ts | 2 +- .../GuardrailViewer/GuardrailViewer.test.tsx | 10 ++++---- .../GuardrailViewer/GuardrailViewer.tsx | 24 +++++++++---------- .../LogDetailContent.integration.test.tsx | 14 +++++------ .../LogDetailsDrawer/LogDetailContent.tsx | 12 +++++----- 18 files changed, 73 insertions(+), 87 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e88de4acd0e..9ba9fd082f3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6021,7 +6021,6 @@ def _get_status_fields( "failure": "guardrail_failed_to_respond", # legacy "guardrail_failed_to_respond": "guardrail_failed_to_respond", # direct "not_run": "not_run", - "skipped": "not_run", } # Set LLM API status diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index dca90e07421..0f5096d0108 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -214,7 +214,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response="no scannable content after message scoping", request_data=data, - guardrail_status="skipped", + guardrail_status="not_run", ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/compliance_checks.py b/litellm/proxy/compliance_checks.py index ef2d8fb6e20..d9cc1d0f4fc 100644 --- a/litellm/proxy/compliance_checks.py +++ b/litellm/proxy/compliance_checks.py @@ -26,7 +26,7 @@ class ComplianceChecker: def __init__(self, data: ComplianceCheckRequest): self.data = data - self.guardrails = [g for g in (data.guardrail_information or []) if g.get("guardrail_status") != "skipped"] + self.guardrails = [g for g in (data.guardrail_information or []) if g.get("guardrail_status") != "not_run"] def _get_guardrails_by_mode(self, mode: str) -> list[dict]: """ diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 651c4bb1963..556b6a4e919 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -42,7 +42,7 @@ if TYPE_CHECKING: router: Final = APIRouter() _EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({}) -_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"skipped": 0, "passed": 1, "flagged": 2, "blocked": 3}) +_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"not_run": 0, "passed": 1, "flagged": 2, "blocked": 3}) _T = TypeVar("_T") @@ -325,7 +325,7 @@ class UsageDetailResponse(BaseModel): class UsageLogEntry(BaseModel): id: str timestamp: str - action: str # blocked | passed | flagged | skipped + action: str # blocked | passed | flagged | not_run score: float | None latency_ms: float | None model: str | None diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 8a131fbfcde..7e11b69108b 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -193,12 +193,12 @@ async def _upsert_rows_with_retry( def guardrail_status_to_action(status: str | None) -> str: - """Map StandardLogging guardrail_status to blocked/passed/flagged/skipped.""" + """Map StandardLogging guardrail_status to blocked/passed/flagged/not_run.""" if not status: return "passed" s: Final = (status or "").lower() - if s == "skipped": - return "skipped" + if s == "not_run": + return "not_run" if "intervened" in s or "block" in s: return "blocked" if "flagged" in s or "fail" in s or "error" in s: @@ -380,7 +380,7 @@ async def process_spend_logs_guardrail_usage( if not isinstance(guardrail_id, str) or not guardrail_id: continue action = guardrail_status_to_action(entry.get("guardrail_status")) - if action != "skipped": + if action != "not_run": key = _MetricsKey(guardrail_id, date_key) daily_guardrail[key]["requests_evaluated"] += 1 if action == "passed": diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cf302b3f27f..1d73542c9bb 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3120,7 +3120,7 @@ class GuardrailMode(TypedDict, total=False): GuardrailStatus = Literal[ - "success", "guardrail_flagged", "guardrail_intervened", "guardrail_failed_to_respond", "not_run", "skipped" + "success", "guardrail_flagged", "guardrail_intervened", "guardrail_failed_to_respond", "not_run" ] # Fields on a guardrail record whose values can quote the caller's prompt: the payload sent to the @@ -3367,7 +3367,6 @@ class StandardLoggingPayloadStatusFields(TypedDict, total=False): - 'guardrail_intervened': Guardrail blocked or modified content - 'guardrail_failed_to_respond': Guardrail had technical failure - 'not_run': No guardrail was run - - 'skipped': Only used per guardrail entry, message scoping left the guardrail nothing to scan """ diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 6e67781944e..70f9bae283b 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -6926,18 +6926,6 @@ def test_get_status_fields_ranks_guardrail_flagged_between_success_and_intervene )["guardrail_status"] == "guardrail_intervened" -def test_get_status_fields_rolls_skipped_entries_up_to_not_run(): - """LIT-6314: a guardrail that message scoping left nothing to scan records a - skipped entry. At request level that means no guardrail ran, and a skipped - entry must never outrank a sibling that did evaluate.""" - skipped = {"guardrail_status": "skipped"} - - assert _get_status_fields("success", [skipped], None)["guardrail_status"] == "not_run" - assert _get_status_fields( - "success", [skipped, {"guardrail_status": "success"}], None - )["guardrail_status"] == "success" - - def test_get_error_information_redacts_provider_key_from_upstream_url(): """A pass-through upstream failure logs the httpx traceback, whose message quotes the upstream URL with the provider key in its query string. That diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index c6d01db45de..4f1163ed806 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1893,7 +1893,7 @@ class TestScanOnlyToolResults: assert data["messages"][4]["content"] == "and then?" -class TestNoScannableContentRecordsSkipped: +class TestNoScannableContentRecordsNotRun: """LIT-6314: a guardrail whose scoping leaves nothing to scan must still persist an evaluation record""" def _system_only_data(self) -> dict: @@ -1904,7 +1904,7 @@ class TestNoScannableContentRecordsSkipped: return metadata.get("standard_logging_guardrail_information") or [] @pytest.mark.asyncio - async def test_skipped_scan_records_skipped_entry(self): + async def test_skipped_scan_records_not_run_entry(self): handler = OpenAIChatCompletionsHandler() guardrail = MockGuardrail(guardrail_name="skip-system-guardrail") guardrail.skip_system_message_in_guardrail = True @@ -1916,7 +1916,7 @@ class TestNoScannableContentRecordsSkipped: entries = self._recorded_entries(data) assert len(entries) == 1 assert entries[0]["guardrail_name"] == "skip-system-guardrail" - assert entries[0]["guardrail_status"] == "skipped" + assert entries[0]["guardrail_status"] == "not_run" @pytest.mark.asyncio async def test_self_recording_guardrail_is_left_alone(self): @@ -1940,7 +1940,7 @@ class TestNoScannableContentRecordsSkipped: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) assert guardrail.last_inputs is not None - assert all(e.get("guardrail_status") != "skipped" for e in self._recorded_entries(data)) + assert all(e.get("guardrail_status") != "not_run" for e in self._recorded_entries(data)) class TestBuildBlockSseChunks: diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index 1df0c1477e2..db87e12ac88 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -685,7 +685,7 @@ async def test_detail_prev_trend_query_is_bounded(): @pytest.mark.asyncio -async def test_logs_report_skipped_entries_as_skipped_not_passed(): +async def test_logs_report_not_run_entries_as_not_run_not_passed(): """LIT-6314: a guardrail that never scanned must not be reported as a pass in the drill-down.""" index_row = MagicMock() index_row.request_id = "req-nr" @@ -697,7 +697,7 @@ async def test_logs_report_skipped_entries_as_skipped_not_passed(): spend_log.startTime = datetime(2026, 4, 22) spend_log.metadata = { "guardrail_information": [ - {"guardrail_name": "db-1", "guardrail_status": "skipped", "duration": 0.0}, + {"guardrail_name": "db-1", "guardrail_status": "not_run", "duration": 0.0}, ] } prisma = _prisma(find_unique=_db_row(), index_find_many=[index_row]) @@ -715,11 +715,11 @@ async def test_logs_report_skipped_entries_as_skipped_not_passed(): end_date=END, user_api_key_dict=ADMIN, ) - assert [log.action for log in resp.logs] == ["skipped"] + assert [log.action for log in resp.logs] == ["not_run"] @pytest.mark.asyncio -async def test_logs_action_passed_filter_excludes_skipped_entries(): +async def test_logs_action_passed_filter_excludes_not_run_entries(): """LIT-6314: filtering the drill-down for passes must not return unscanned requests.""" index_row = MagicMock() index_row.request_id = "req-nr" @@ -729,7 +729,7 @@ async def test_logs_action_passed_filter_excludes_skipped_entries(): spend_log.request_id = "req-nr" spend_log.model = "gpt-4o-mini" spend_log.startTime = datetime(2026, 4, 22) - spend_log.metadata = {"guardrail_information": [{"guardrail_name": "db-1", "guardrail_status": "skipped"}]} + spend_log.metadata = {"guardrail_information": [{"guardrail_name": "db-1", "guardrail_status": "not_run"}]} prisma = _prisma(find_unique=_db_row(), index_find_many=[index_row]) prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[spend_log]) handler = _config_handler() diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index cb85eb8b5ec..69ec098b840 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -350,15 +350,15 @@ async def test_zero_and_non_int_usage_counters_are_skipped(): @pytest.mark.asyncio -async def test_skipped_entries_are_indexed_but_not_counted_as_evaluations(): +async def test_not_run_entries_are_indexed_but_not_counted_as_evaluations(): """ - LIT-6314 records a skipped entry when message scoping leaves a guardrail + LIT-6314 records a not_run entry when message scoping leaves a guardrail nothing to scan. The guardrail never evaluated the request, so counting it as a passed evaluation would inflate daily pass rates; it still gets an index row so per-request drill-down finds the spend log. """ prisma = _prisma() - logs = [_payload("r1", guardrail_status="skipped"), _payload("r2")] + logs = [_payload("r1", guardrail_status="not_run"), _payload("r2")] await process_spend_logs_guardrail_usage(prisma, logs) @@ -370,26 +370,26 @@ async def test_skipped_entries_are_indexed_but_not_counted_as_evaluations(): @pytest.mark.asyncio -async def test_skipped_entry_shares_index_key_with_evaluated_sibling_of_same_name(): +async def test_not_run_entry_shares_index_key_with_evaluated_sibling_of_same_name(): """ - The skipped entry from the shared base guardrail carries only guardrail_name, + The not_run entry from the shared base guardrail carries only guardrail_name, while the evaluated entry from the same guardrail (e.g. content filter on the output of a logging_only run) carries its guardrail_id. Keying them differently - lists one request twice in the monitor, once as skipped and once as passed. + lists one request twice in the monitor, once as not_run and once as passed. """ prisma = _prisma() payload = _payload("r1") payload["metadata"] = json.dumps( { "guardrail_information": [ - {"guardrail_name": "cf", "guardrail_status": "skipped"}, + {"guardrail_name": "cf", "guardrail_status": "not_run"}, { "guardrail_name": "cf", "guardrail_id": "cf-uuid", "policy_id": "pol-1", "guardrail_status": "success", }, - {"guardrail_name": "other", "guardrail_status": "skipped"}, + {"guardrail_name": "other", "guardrail_status": "not_run"}, ] } ) @@ -406,7 +406,7 @@ async def test_skipped_entry_shares_index_key_with_evaluated_sibling_of_same_nam @pytest.mark.asyncio -async def test_malformed_skipped_entry_does_not_drop_the_batch(): +async def test_malformed_not_run_entry_does_not_drop_the_batch(): prisma = _prisma() payload = _payload("r1") payload["metadata"] = json.dumps( @@ -428,10 +428,10 @@ async def test_malformed_skipped_entry_does_not_drop_the_batch(): @pytest.mark.asyncio -async def test_batch_of_only_skipped_entries_writes_no_metrics_row(): +async def test_batch_of_only_not_run_entries_writes_no_metrics_row(): prisma = _prisma() - await process_spend_logs_guardrail_usage(prisma, [_payload("r1", guardrail_status="skipped")]) + await process_spend_logs_guardrail_usage(prisma, [_payload("r1", guardrail_status="not_run")]) assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 0 index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py index 9587e3d95ee..8382a5ada96 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_compliance_endpoints.py @@ -591,17 +591,17 @@ class TestModeMatching: assert mode in _guaranteed_modes(g_mode), (g_mode, mode) -class TestSkippedGuardrails: - """LIT-6314 logs a skipped entry for a guardrail that message scoping left nothing to scan.""" +class TestNotRunGuardrails: + """LIT-6314 logs a not_run entry for a guardrail that message scoping left nothing to scan.""" - def test_skipped_alone_never_evidences_compliance(self): + def test_not_run_alone_never_evidences_compliance(self): data = ComplianceCheckRequest( request_id="req-601", user_id="user-1", model="gpt-4", timestamp="2026-02-17T00:00:00Z", guardrail_information=[ - {"guardrail_name": "pii_detection", "guardrail_status": "skipped", "guardrail_mode": "pre_call"}, + {"guardrail_name": "pii_detection", "guardrail_status": "not_run", "guardrail_mode": "pre_call"}, ], ) results = {c.check_name: c.passed for c in ComplianceChecker(data).check_eu_ai_act()} @@ -609,7 +609,7 @@ class TestSkippedGuardrails: assert results["Content screened before LLM"] is False assert results["Audit record complete"] is False - def test_skipped_sibling_does_not_fail_a_passing_request(self): + def test_not_run_sibling_does_not_fail_a_passing_request(self): data = ComplianceCheckRequest( request_id="req-602", user_id="user-1", @@ -618,7 +618,7 @@ class TestSkippedGuardrails: pii_detected=True, guardrail_information=[ {"guardrail_name": "pii_detection", "guardrail_status": "success", "guardrail_mode": "pre_call"}, - {"guardrail_name": "system_only", "guardrail_status": "skipped", "guardrail_mode": "pre_call"}, + {"guardrail_name": "system_only", "guardrail_status": "not_run", "guardrail_mode": "pre_call"}, ], ) results = {c.check_name: c.passed for c in ComplianceChecker(data).check_gdpr()} diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx index af5284830dc..083b1e5f3e2 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx @@ -96,14 +96,14 @@ describe("GuardrailsMonitor LogViewer drawer", () => { }); }); -describe("GuardrailsMonitor LogViewer skipped rows", () => { - it("renders a skipped log as a neutral Skipped badge instead of a pass or failure", () => { +describe("GuardrailsMonitor LogViewer not_run rows", () => { + it("renders a not_run log as a neutral Not run badge instead of a pass or failure", () => { renderWithProviders( - , + , ); const row = screen.getByRole("button", { name: /system prompt only/ }); - expect(within(row).getByText("Skipped")).toHaveClass("text-muted-foreground"); + expect(within(row).getByText("Not run")).toHaveClass("text-muted-foreground"); expect(within(row).queryByText("Passed")).not.toBeInTheDocument(); expect(within(row).queryByText("Blocked")).not.toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx index b113c49e0a3..2abd699ba86 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx @@ -10,15 +10,15 @@ import type { LogEntry as ViewLogsLogEntry } from "@/components/view_logs/column import type { LogEntry } from "./mockData"; const actionConfig: Record< - "blocked" | "passed" | "flagged" | "skipped", + "blocked" | "passed" | "flagged" | "not_run", { icon: React.ElementType; color: string; bg: string; border: string; label: string } > = { - skipped: { + not_run: { icon: MinusCircle, color: "text-muted-foreground", bg: "bg-muted", border: "border-border", - label: "Skipped", + label: "Not run", }, blocked: { icon: X, diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts b/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts index 7053efcf88d..591d5cd3edd 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/mockData.ts @@ -10,7 +10,7 @@ export interface LogEntry { input_snippet?: string; output_snippet?: string; score?: number; - action: "blocked" | "passed" | "flagged" | "skipped"; + action: "blocked" | "passed" | "flagged" | "not_run"; model?: string; reason?: string; latency_ms?: number; diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx index 740ef857690..7f343211596 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx @@ -16,7 +16,7 @@ const PresidioPath = "@/components/view_logs/GuardrailViewer/PresidioDetectedEnt const BedrockPath = "@/components/view_logs/GuardrailViewer/BedrockGuardrailDetails"; const skippedPreCall: Partial = { - guardrail_status: "skipped", + guardrail_status: "not_run", guardrail_mode: "pre_call", guardrail_response: "no scannable content after message scoping", start_time: null, @@ -68,15 +68,15 @@ describe("GuardrailViewer", () => { expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); }); - it("renders skipped as SKIPPED (muted) and keeps it out of the evaluated and passed counts", async () => { + it("renders not_run as NOT RUN (muted) and keeps it out of the evaluated and passed counts", async () => { const user = userEvent.setup(); const data = makeGuardrailInformation(skippedPreCall); renderWithProviders(); expect(screen.getByText(/0 guardrails evaluated/)).toBeInTheDocument(); expect(screen.getByText(/0 Passed/)).toHaveClass("text-muted-foreground"); - expect(screen.getByText(/1 Skipped/)).toBeInTheDocument(); - const badge = screen.getByText("SKIPPED"); + expect(screen.getByText(/1 Not run/)).toBeInTheDocument(); + const badge = screen.getByText("NOT RUN"); expect(badge).toHaveClass("text-muted-foreground"); expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); expect(screen.queryByText(/^T\+/)).not.toBeInTheDocument(); @@ -85,7 +85,7 @@ describe("GuardrailViewer", () => { expect(screen.getByText("no scannable content after message scoping")).toBeInTheDocument(); }); - it("anchors the lifecycle timeline on timed entries when an untimed skipped entry sorts first", () => { + it("anchors the lifecycle timeline on timed entries when an untimed not_run entry sorts first", () => { const skipped = makeGuardrailInformation({ ...skippedPreCall, guardrail_name: "skipped-rail" }); const ran = makeGuardrailInformation(ranPostCall); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx index c67d682a233..1de0e3878b2 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx @@ -134,13 +134,13 @@ const getTotalMasked = (entry: GuardrailInformation): number => { ); }; -type EntryOutcome = "passed" | "flagged" | "failed" | "skipped"; +type EntryOutcome = "passed" | "flagged" | "failed" | "not_run"; const getEntryOutcome = (entry: GuardrailInformation): EntryOutcome => { const status = (entry.guardrail_status ?? "").toLowerCase(); if (status === "success") return "passed"; if (status === "guardrail_flagged") return "flagged"; - if (status === "skipped") return "skipped"; + if (status === "not_run") return "not_run"; return "failed"; }; @@ -150,18 +150,18 @@ const OUTCOME_LABEL: Record = { passed: "PASSED", flagged: "FLAGGED", failed: "FAILED", - skipped: "SKIPPED", + not_run: "NOT RUN", }; const OUTCOME_BADGE_CLASS: Record = { passed: "bg-success/15 text-success border border-success/20", flagged: "bg-warning/15 text-warning border border-warning/20", failed: "bg-destructive/15 text-destructive border border-destructive/20", - skipped: "bg-muted text-muted-foreground border border-border", + not_run: "bg-muted text-muted-foreground border border-border", }; const getHeaderOutcome = (counts: { evaluated: number; passed: number; flagged: number }): EntryOutcome => { - if (counts.evaluated === 0) return "skipped"; + if (counts.evaluated === 0) return "not_run"; if (counts.passed === counts.evaluated) return "passed"; if (counts.passed + counts.flagged === counts.evaluated) return "flagged"; return "failed"; @@ -242,7 +242,7 @@ const FlagCircleIcon = ({ className }: { className?: string }) => ( const OutcomeIcon = ({ outcome }: { outcome: EntryOutcome }) => { if (outcome === "passed") return ; if (outcome === "flagged") return ; - if (outcome === "skipped") return ; + if (outcome === "not_run") return ; return ; }; @@ -675,7 +675,7 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { )} - {outcome === "skipped" && typeof guardrailResponse === "string" && ( + {outcome === "not_run" && typeof guardrailResponse === "string" && (

{guardrailResponse}

)} @@ -717,8 +717,8 @@ const GuardrailViewer = ({ data, accessToken, logEntry }: GuardrailViewerProps) const passedCount = guardrailEntries.filter(isEntrySuccess).length; const flaggedCount = guardrailEntries.filter((e) => getEntryOutcome(e) === "flagged").length; - const skippedCount = guardrailEntries.filter((e) => getEntryOutcome(e) === "skipped").length; - const evaluatedCount = guardrailEntries.length - skippedCount; + const notRunCount = guardrailEntries.filter((e) => getEntryOutcome(e) === "not_run").length; + const evaluatedCount = guardrailEntries.length - notRunCount; const allPassed = evaluatedCount > 0 && passedCount === evaluatedCount; const headerOutcome = getHeaderOutcome({ evaluated: evaluatedCount, passed: passedCount, flagged: flaggedCount }); @@ -778,11 +778,11 @@ const GuardrailViewer = ({ data, accessToken, logEntry }: GuardrailViewerProps) {flaggedCount} Flagged )} - {skippedCount > 0 && ( + {notRunCount > 0 && ( - {skippedCount} Skipped + {notRunCount} Not run )} diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx index cbc377c1b73..3637b6c55a3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.integration.test.tsx @@ -637,20 +637,20 @@ describe("GuardrailJumpLink", () => { }); it.each([ - [["success", "skipped"], "text-success", "\u2713"], - [["guardrail_intervened", "skipped"], "text-destructive", "\u2717"], - ])("ignores skipped when styling %j as %s", (statuses, expectedClass, glyph) => { + [["success", "not_run"], "text-success", "\u2713"], + [["guardrail_intervened", "not_run"], "text-destructive", "\u2717"], + ])("ignores not_run when styling %j as %s", (statuses, expectedClass, glyph) => { render( ({ guardrail_status: s }))} />); - const pill = screen.getByText(/1 guardrail evaluated, 1 skipped/); + const pill = screen.getByText(/1 guardrail evaluated, 1 not run/); expect(pill).toHaveClass(expectedClass); expect(pill).toHaveTextContent(glyph); }); - it("renders an all skipped request as neutral rather than passed", () => { - render(); + it("renders an all not_run request as neutral rather than passed", () => { + render(); - const pill = screen.getByText(/0 guardrails evaluated, 1 skipped/); + const pill = screen.getByText(/0 guardrails evaluated, 1 not run/); expect(pill).toHaveClass("text-muted-foreground"); expect(pill).not.toHaveTextContent("\u2713"); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx index e052faa2228..30cd96935f7 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx @@ -700,15 +700,15 @@ const GUARDRAIL_JUMP_LINK_STYLE = { passed: { className: "border border-success/20 bg-success/10 text-success", glyph: "\u2713" }, flagged: { className: "border border-warning/20 bg-warning/10 text-warning", glyph: "\u26A0" }, failed: { className: "border border-destructive/20 bg-destructive/10 text-destructive", glyph: "\u2717" }, - skipped: { className: "border border-border bg-muted text-muted-foreground", glyph: "\u2013" }, + not_run: { className: "border border-border bg-muted text-muted-foreground", glyph: "\u2013" }, } as const; const isPassedStatus = (status: unknown) => status === "pass" || status === "passed" || status === "success"; const isFlaggedStatus = (status: unknown) => status === "flagged" || status === "guardrail_flagged"; -const isSkippedStatus = (status: unknown) => status === "skipped"; +const isNotRunStatus = (status: unknown) => status === "not_run"; const guardrailJumpLinkOutcome = (evaluated: unknown[]): keyof typeof GUARDRAIL_JUMP_LINK_STYLE => { - if (evaluated.length === 0) return "skipped"; + if (evaluated.length === 0) return "not_run"; if (evaluated.every(isPassedStatus)) return "passed"; if (evaluated.every((s) => isPassedStatus(s) || isFlaggedStatus(s))) return "flagged"; return "failed"; @@ -716,8 +716,8 @@ const guardrailJumpLinkOutcome = (evaluated: unknown[]): keyof typeof GUARDRAIL_ export function GuardrailJumpLink({ guardrailEntries }: { guardrailEntries: any[] }) { const statuses = guardrailEntries.map((e) => e?.guardrail_status || e?.status); - const evaluated = statuses.filter((s) => !isSkippedStatus(s)); - const skippedCount = statuses.length - evaluated.length; + const evaluated = statuses.filter((s) => !isNotRunStatus(s)); + const notRunCount = statuses.length - evaluated.length; const { className, glyph } = GUARDRAIL_JUMP_LINK_STYLE[guardrailJumpLinkOutcome(evaluated)]; const handleClick = () => { @@ -743,7 +743,7 @@ export function GuardrailJumpLink({ guardrailEntries }: { guardrailEntries: any[ > {glyph} {evaluated.length} guardrail {evaluated.length !== 1 ? "s" : ""} evaluated - {skippedCount > 0 ? `, ${skippedCount} skipped` : ""} + {notRunCount > 0 ? `, ${notRunCount} not run` : ""} {"\u2193"} From 0519d8634600bc404afc48a5d111d7869644e00e Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 23:49:03 +0000 Subject: [PATCH 063/107] fix(guardrails): stop labelling image-only input as a not_run scoping skip Images without text were never dispatched to guardrails before this change, so that gap is not a message scoping skip and must not get a not_run entry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/guardrail_translation/handler.py | 2 +- .../test_openai_guardrail_handler.py | 21 +++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 0f5096d0108..9d790aa71d4 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -210,7 +210,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings=tool_call_task_mappings, ) - elif not guardrail_to_apply.records_own_guardrail_information: + elif not images_to_check and not guardrail_to_apply.records_own_guardrail_information: guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response="no scannable content after message scoping", request_data=data, diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 4f1163ed806..5c971fe2c90 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1942,6 +1942,27 @@ class TestNoScannableContentRecordsNotRun: assert guardrail.last_inputs is not None assert all(e.get("guardrail_status") != "not_run" for e in self._recorded_entries(data)) + @pytest.mark.asyncio + async def test_image_only_content_is_not_reported_as_not_run(self): + """Images are only scanned alongside text, so an image-only request is a + pre-existing scan gap, not a message-scoping skip, and must not be labelled one""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="image-guardrail") + guardrail.skip_system_message_in_guardrail = True + data = { + "messages": [ + {"role": "system", "content": "SYSTEM-PROMPT"}, + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}], + }, + ] + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert self._recorded_entries(data) == [] + class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks""" From 52da64a45b76f516b2a0354af93a2aa5b851eec3 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 23:56:32 +0000 Subject: [PATCH 064/107] fix(guardrails): only cite message scoping in the not_run reason when scoping is on A request whose messages carry no scannable content at all, with no skip flag set, now records the neutral reason no scannable content instead of blaming configuration Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../openai/chat/guardrail_translation/handler.py | 6 +++++- .../test_openai_guardrail_handler.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 9d790aa71d4..0be4b6a3a20 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -212,7 +212,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation): elif not images_to_check and not guardrail_to_apply.records_own_guardrail_information: guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response="no scannable content after message scoping", + guardrail_json_response=( + "no scannable content after message scoping" + if skip_system or skip_tool or scan_only_tool_results + else "no scannable content" + ), request_data=data, guardrail_status="not_run", ) diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5c971fe2c90..40dc5e2df2c 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1917,6 +1917,21 @@ class TestNoScannableContentRecordsNotRun: assert len(entries) == 1 assert entries[0]["guardrail_name"] == "skip-system-guardrail" assert entries[0]["guardrail_status"] == "not_run" + assert entries[0]["guardrail_response"] == "no scannable content after message scoping" + + @pytest.mark.asyncio + async def test_empty_content_without_scoping_does_not_blame_scoping(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="unscoped-guardrail") + data = {"messages": [{"role": "user", "content": None}]} + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.last_inputs is None + entries = self._recorded_entries(data) + assert len(entries) == 1 + assert entries[0]["guardrail_status"] == "not_run" + assert entries[0]["guardrail_response"] == "no scannable content" @pytest.mark.asyncio async def test_self_recording_guardrail_is_left_alone(self): From 5ceec4c21e9dd35f7b0693bdef327a1e59d7d3c8 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 15 Sep 2026 00:03:02 +0000 Subject: [PATCH 065/107] fix(guardrails): cite message scoping only when an unscoped pass finds content The not_run reason now says after message scoping only when the same messages carry text or tool calls without the skip flags applied. A request that is empty to begin with, whatever the flags, records no scannable content Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../openai/chat/guardrail_translation/handler.py | 14 +++++++++++++- .../test_openai_guardrail_handler.py | 4 +++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 0be4b6a3a20..e8179e9921e 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -211,10 +211,22 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) elif not images_to_check and not guardrail_to_apply.records_own_guardrail_information: + unscoped_texts: Final[list[str]] = [] + unscoped_tool_calls: Final[list[ChatCompletionToolParam]] = [] + for unscoped_idx, unscoped_message in enumerate(messages): + self._extract_inputs( + message=unscoped_message, + msg_idx=unscoped_idx, + texts_to_check=unscoped_texts, + images_to_check=[], + tool_calls_to_check=unscoped_tool_calls, + text_task_mappings=[], + tool_call_task_mappings=[], + ) guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=( "no scannable content after message scoping" - if skip_system or skip_tool or scan_only_tool_results + if unscoped_texts or unscoped_tool_calls else "no scannable content" ), request_data=data, diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 40dc5e2df2c..754c89d3d20 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1920,9 +1920,11 @@ class TestNoScannableContentRecordsNotRun: assert entries[0]["guardrail_response"] == "no scannable content after message scoping" @pytest.mark.asyncio - async def test_empty_content_without_scoping_does_not_blame_scoping(self): + @pytest.mark.parametrize("skip_system", [False, True]) + async def test_empty_content_does_not_blame_scoping(self, skip_system: bool): handler = OpenAIChatCompletionsHandler() guardrail = MockGuardrail(guardrail_name="unscoped-guardrail") + guardrail.skip_system_message_in_guardrail = skip_system data = {"messages": [{"role": "user", "content": None}]} await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) From 16c326537f5aa06c597fa198a5cb9e7a01fd9327 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:39:54 -0700 Subject: [PATCH 066/107] fix(guardrails): define UnappliableRequestRewrite in the shared guardrail translation utils The three guardrail translation handlers imported the exception from the proxy policy engine through a function-local import, which CodeQL flagged as a cyclic import. The exception and its helper now live next to the handlers in the shared guardrail translation utils, and the tests import it from there. The Prompt Security modify-mode helper is also restructured into early-return TypedDict displays so the LIT002 budget stays at its limit --- .../base_llm/guardrail_translation/utils.py | 11 ++++++++-- .../prompt_security/prompt_security.py | 21 ++++++++++++++----- .../proxy/policy_engine/pipeline_executor.py | 9 -------- .../test_anthropic_guardrail_handler.py | 2 +- .../test_openai_guardrail_handler.py | 2 +- ...test_openai_responses_guardrail_handler.py | 4 ++-- .../guardrail_hooks/test_crowdstrike_aidr.py | 2 +- 7 files changed, 30 insertions(+), 21 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index a80e20c5404..c47a56b5ea9 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -413,7 +413,14 @@ def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped -def unappliable_request_rewrite(guardrail_name: str | None) -> Exception: - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite +class UnappliableRequestRewrite(Exception): + def __init__(self, guardrail_name: str) -> None: + super().__init__( + f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, " + "so the request was rejected rather than sent unrewritten" + ) + self.guardrail_name: Final = guardrail_name + +def unappliable_request_rewrite(guardrail_name: str | None) -> UnappliableRequestRewrite: return UnappliableRequestRewrite(guardrail_name or "unknown") diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 4581b8b863f..3f29b3e751c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -38,16 +38,27 @@ class PromptSecurityGuardrailMissingSecrets(Exception): pass +def _inputs_with_structured_messages( + inputs: GenericGuardrailAPIInputs, rewritten_messages: Sequence[AllMessageValues] | None +) -> GenericGuardrailAPIInputs: + if rewritten_messages is None: + return inputs + patched: Final[GenericGuardrailAPIInputs] = { + **inputs, + "structured_messages": list(rewritten_messages), # mutable-ok: the TypedDict field is declared as a list + } + return patched + + def _inputs_with_modifications( inputs: GenericGuardrailAPIInputs, modified_texts: list[str], rewritten_messages: Sequence[AllMessageValues] | None, ) -> GenericGuardrailAPIInputs: - texts_patch: Final[GenericGuardrailAPIInputs] = {"texts": modified_texts} if modified_texts else {} - messages_patch: Final[GenericGuardrailAPIInputs] = ( - {"structured_messages": list(rewritten_messages)} if rewritten_messages is not None else {} - ) - return {**inputs, **texts_patch, **messages_patch} + if not modified_texts: + return _inputs_with_structured_messages(inputs, rewritten_messages) + with_texts: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": modified_texts} + return _inputs_with_structured_messages(with_texts, rewritten_messages) class _ProtectVerdict(TypedDict, total=False): diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index ad45781d5d2..ed193c7f434 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -58,15 +58,6 @@ class UndeliverableStreamRewrite(Exception): self.guardrail_name: Final = guardrail_name -class UnappliableRequestRewrite(Exception): - def __init__(self, guardrail_name: str) -> None: - super().__init__( - f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, " - "so the request was rejected rather than sent unrewritten" - ) - self.guardrail_name: Final = guardrail_name - - def _tool_call_shape(tool_call: object) -> tuple[object, object]: plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call function: Final = plain.get("function") if isinstance(plain, Mapping) else None diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index d6fd30638cc..0882b329c49 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2296,7 +2296,7 @@ class TestPerMessageTextWriteBack: @pytest.mark.asyncio async def test_one_text_per_row_over_a_system_prompt_is_rejected_by_name(self): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite data = { "model": "claude-sonnet-4-5", diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 4e2291fdec1..8ee0e982aa0 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1917,7 +1917,7 @@ class TestPerMessageTextWriteBack: @pytest.mark.asyncio async def test_fewer_texts_than_extracted_over_a_tool_message_is_rejected(self): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite handler = OpenAIChatCompletionsHandler() original_messages = [ 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 ac719da169c..48d86384633 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 @@ -2426,7 +2426,7 @@ class TestPerMessageRewriteWriteBack: @pytest.mark.asyncio async def test_texts_only_per_message_answer_is_rejected_by_name(self): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite guardrail = _per_message_redactor() data = _tool_replay_request() @@ -2453,7 +2453,7 @@ class TestPerMessageRewriteWriteBack: @pytest.mark.asyncio async def test_texts_only_per_message_answer_over_a_string_input_is_rejected_by_name(self): - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite guardrail = _per_message_redactor() data = _string_input_request() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index a9ca13a463d..9849ad7ec88 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -1820,7 +1820,7 @@ async def test_unalignable_rewrite_is_rejected_never_sent_unredacted( Skipping the write-back would hand the model the unredacted text, so a guardrail could be bypassed by adding ``instructions`` or a tool call. """ - from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite data: dict[str, object] = {"model": "gpt-4o", "input": responses_input} if instructions is not None: From 8d5d5d58245d5765f8b5fd13d234d952a97e2030 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 19:03:55 -0700 Subject: [PATCH 067/107] fix(bedrock): grant rerank, retrieve, agent, and agentcore actions in the web identity session policy --- litellm/llms/bedrock/base_aws_llm.py | 121 +++++++++++------- .../test_web_identity_session_policy.py | 57 ++++++++- 2 files changed, 124 insertions(+), 54 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index f52c1cec6a8..70869f69e3a 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -11,6 +11,7 @@ from concurrent.futures import ThreadPoolExecutor from datetime import datetime from functools import partial from threading import Lock +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload import httpx @@ -96,6 +97,76 @@ def _assume_role_params( ) +_SecureTransportBool = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]}) + + +class _SecureTransportCondition(TypedDict): + Bool: ReadOnly[_SecureTransportBool] + + +class _SessionPolicyStatement(TypedDict): + Sid: ReadOnly[str] + Effect: ReadOnly[Literal["Allow"]] + Action: ReadOnly[tuple[str, ...]] + Resource: ReadOnly[Literal["*"]] + Condition: ReadOnly[_SecureTransportCondition] + + +class WebIdentitySessionPolicy(TypedDict): + Version: ReadOnly[Literal["2012-10-17"]] + Statement: ReadOnly[tuple[_SessionPolicyStatement, ...]] + + +_WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( + { + "BedrockLiteLLM": ( + "bedrock:InvokeModel", + "bedrock:InvokeModelWithResponseStream", + "bedrock:CountTokens", + "bedrock:Rerank", + "bedrock:Retrieve", + "bedrock:ListKnowledgeBases", + "bedrock:InvokeAgent", + "bedrock:ApplyGuardrail", + "bedrock:GetGuardrail", + "bedrock:ListGuardrails", + ), + "BedrockAgentCoreLiteLLM": ( + "bedrock-agentcore:InvokeAgentRuntime", + "bedrock-agentcore:InvokeGateway", + ), + "ClaudePlatformLiteLLM": ( + "aws-external-anthropic:CreateInference", + "aws-external-anthropic:CreateBatchInference", + "aws-external-anthropic:CancelBatchInference", + "aws-external-anthropic:DeleteBatchInference", + "aws-external-anthropic:CountTokens", + "aws-external-anthropic:Get*", + "aws-external-anthropic:List*", + ), + "BedrockMantleLiteLLM": ("bedrock-mantle:CreateInference",), + } +) + +_SECURE_TRANSPORT_ONLY: Final = _SecureTransportCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"})) + + +def build_web_identity_session_policy() -> WebIdentitySessionPolicy: + return WebIdentitySessionPolicy( + Version="2012-10-17", + Statement=tuple( + _SessionPolicyStatement( + Sid=sid, + Effect="Allow", + Action=actions, + Resource="*", + Condition=_SECURE_TRANSPORT_ONLY, + ) + for sid, actions in _WEB_IDENTITY_SESSION_POLICY_ACTIONS.items() + ), + ) + + class BedrockRequestTarget(BaseModel): aws_region_name: str aws_bedrock_runtime_endpoint: str | None @@ -940,60 +1011,12 @@ class BaseAWSLLM(SignsRequestsWithAWS): # auth only (static creds + IRSA take other code paths). # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html - bedrock_session_policy: Final = { - "Version": "2012-10-17", - "Statement": [ - { - "Sid": "BedrockLiteLLM", - "Effect": "Allow", - "Action": [ - "bedrock:InvokeModel", - "bedrock:InvokeModelWithResponseStream", - "bedrock:CountTokens", - "bedrock:ApplyGuardrail", - "bedrock:GetGuardrail", - "bedrock:ListGuardrails", - ], - "Resource": "*", - "Condition": {"Bool": {"aws:SecureTransport": "true"}}, - }, - # Claude Platform on AWS (added by #27678 for the - # ``bedrock/claude_platform/`` route) lives under - # a separate IAM action namespace; without these entries - # the OIDC path 403s on every claude_platform request - # even with a fully permissive identity policy (#30200). - { - "Sid": "ClaudePlatformLiteLLM", - "Effect": "Allow", - "Action": [ - "aws-external-anthropic:CreateInference", - "aws-external-anthropic:CreateBatchInference", - "aws-external-anthropic:CancelBatchInference", - "aws-external-anthropic:DeleteBatchInference", - "aws-external-anthropic:CountTokens", - "aws-external-anthropic:Get*", - "aws-external-anthropic:List*", - ], - "Resource": "*", - "Condition": {"Bool": {"aws:SecureTransport": "true"}}, - }, - { - "Sid": "BedrockMantleLiteLLM", - "Effect": "Allow", - "Action": [ - "bedrock-mantle:CreateInference", - ], - "Resource": "*", - "Condition": {"Bool": {"aws:SecureTransport": "true"}}, - }, - ], - } assume_role_params: Final = { "RoleArn": aws_role_name, "RoleSessionName": aws_session_name, "WebIdentityToken": oidc_token, "DurationSeconds": 3600, - "Policy": json.dumps(bedrock_session_policy, separators=(",", ":")), + "Policy": json.dumps(build_web_identity_session_policy(), separators=(",", ":")), } # Add ExternalId parameter if provided diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index 3ea840519f9..cd5b14c8d37 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -32,6 +32,8 @@ action. import base64 import json from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -49,9 +51,9 @@ _CLAUDE_PLATFORM_ACTIONS = { } -def _captured_policy() -> dict: - """Run _auth_with_web_identity_token under mocks + return the parsed - Policy dict that was actually sent to STS.""" +def _captured_policy_document() -> str: + """Run _auth_with_web_identity_token under mocks + return the Policy + JSON document that was actually sent to STS.""" from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM base = BaseAWSLLM() @@ -84,8 +86,15 @@ def _captured_policy() -> dict: mock_sts.assume_role_with_web_identity.assert_called_once() kwargs = mock_sts.assume_role_with_web_identity.call_args.kwargs - policy_str = kwargs["Policy"] - return json.loads(policy_str) + return kwargs["Policy"] + + +def _captured_policy() -> dict: + return json.loads(_captured_policy_document()) + + +def _granted_actions(policy: dict) -> frozenset[str]: + return frozenset(action for stmt in policy["Statement"] for action in stmt["Action"]) def _statement_by_sid(policy: dict, sid: str) -> dict: @@ -308,3 +317,41 @@ class TestPolicyTransportConditions: "ClaudePlatformLiteLLM must require aws:SecureTransport=true " "to keep parity with the bedrock statement" ) + + +_STS_SESSION_POLICY_PLAINTEXT_LIMIT: Final = 2048 + +_BEDROCK_ROUTE_ACTIONS: Final = MappingProxyType( + { + "model/{model_id}/invoke": "bedrock:InvokeModel", + "model/{model_id}/invoke-with-response-stream": "bedrock:InvokeModelWithResponseStream", + "model/{model_id}/converse": "bedrock:InvokeModel", + "model/{model_id}/converse-stream": "bedrock:InvokeModelWithResponseStream", + "model/{model_id}/count-tokens": "bedrock:CountTokens", + "guardrail/{guardrail_id}/version/{version}/apply": "bedrock:ApplyGuardrail", + "rerank": "bedrock:Rerank", + "knowledgebases/{knowledge_base_id}/retrieve": "bedrock:Retrieve", + "knowledgebases": "bedrock:ListKnowledgeBases", + "agents/{agent_id}/agentAliases/{alias_id}/sessions/{session_id}/text": "bedrock:InvokeAgent", + "runtimes/{agent_runtime_arn}/invocations": "bedrock-agentcore:InvokeAgentRuntime", + "mcp": "bedrock-agentcore:InvokeGateway", + } +) + + +class TestSessionPolicyGrantsEveryBedrockRoute: + """LIT-7348: ``/rerank`` authorizes against ``bedrock:Rerank``, which the + ceiling never granted, so rerank 403d on web identity auth while static + credentials and IRSA worked. Each route the bedrock package signs with the + web identity session maps to the IAM action it authorizes against, and the + ceiling must grant every one of them.""" + + @pytest.mark.parametrize(("route", "action"), sorted(_BEDROCK_ROUTE_ACTIONS.items())) + def test_route_action_is_granted_by_the_ceiling(self, route: str, action: str): + assert action in _granted_actions(_captured_policy()), ( + f"/{route} authorizes against {action}, which the session policy does not grant, " + "so it 403s on web identity auth" + ) + + def test_policy_document_fits_the_sts_plaintext_limit(self): + assert len(_captured_policy_document()) <= _STS_SESSION_POLICY_PLAINTEXT_LIMIT From 7255a201a341b6e15f68a5a02a7abb0d3f6ba759 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:19:15 +0000 Subject: [PATCH 068/107] fix(health): resolve stored credentials for realtime checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/realtime_api/main.py | 38 +++++++++++++------- tests/test_litellm/realtime_api/test_main.py | 32 +++++++++++++++++ 2 files changed, 57 insertions(+), 13 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 44c47af57f4..59aa63eb8a5 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -28,7 +28,7 @@ from litellm.types.realtime import ( ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, LlmProviders -from litellm.utils import ProviderConfigManager +from litellm.utils import ProviderConfigManager, load_credentials_from_list from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -56,6 +56,12 @@ base_llm_http_handler = BaseLLMHTTPHandler() _EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({}) +def _model_params_with_stored_credentials(model_params: Mapping[str, Any]) -> Mapping[str, Any]: + hydrated: Final = dict(model_params) + load_credentials_from_list(hydrated) + return MappingProxyType(hydrated) + + def _with_resolved_session_model(session: dict[str, object], model_name: str) -> dict[str, object]: if "model" not in session: return session @@ -629,34 +635,40 @@ async def _realtime_health_check( """ import websockets + resolved_params: Final = _model_params_with_stored_credentials(model_params or _EMPTY_MODEL_PARAMS) + resolved_api_key: Final = cast(str | None, api_key or resolved_params.get("api_key")) + resolved_api_base: Final = cast(str | None, api_base or resolved_params.get("api_base")) + resolved_api_version: Final = cast(str | None, api_version or resolved_params.get("api_version")) url: str | None = None auth_headers: Final = _realtime_health_check_auth_headers( custom_llm_provider=custom_llm_provider, - api_key=api_key, - model_params=model_params or _EMPTY_MODEL_PARAMS, + api_key=resolved_api_key, + model_params=resolved_params, ) if custom_llm_provider == "azure": resolved_protocol, azure_query_params = _azure_realtime_health_protocol( model=model, realtime_protocol=realtime_protocol, - model_params=model_params or _EMPTY_MODEL_PARAMS, + model_params=resolved_params, ) url = azure_realtime._construct_url( - api_base=api_base or "", + api_base=resolved_api_base or "", model=model, - api_version=api_version or "2024-10-01-preview", + api_version=resolved_api_version or "2024-10-01-preview", realtime_protocol=resolved_protocol, query_params=azure_query_params, ) elif custom_llm_provider == "openai": url = openai_realtime._construct_url( - api_base=api_base or "https://api.openai.com/", + api_base=resolved_api_base or "https://api.openai.com/", query_params={"model": model}, ) elif custom_llm_provider == "xai": - url = xai_realtime._construct_url(api_base=api_base or "https://api.x.ai/v1", query_params={"model": model}) + url = xai_realtime._construct_url( + api_base=resolved_api_base or "https://api.x.ai/v1", query_params={"model": model} + ) elif custom_llm_provider == "vertex_ai": - vertex_model_params: Final = model_params or {} + vertex_model_params: Final = dict(resolved_params) resolved_location: Final = vertex_llm_base.get_vertex_region( vertex_region=VertexBase.safe_get_vertex_ai_location(vertex_model_params), model=model, @@ -675,19 +687,19 @@ async def _realtime_health_check( project=resolved_project, location=resolved_location, ) - url = vertex_realtime_config.get_complete_url(api_base=api_base, model=model) - ssl_context = get_shared_realtime_ssl_context() + url = vertex_realtime_config.get_complete_url(api_base=resolved_api_base, model=model) + vertex_ssl_context: Final = get_shared_realtime_ssl_context() headers: Final = vertex_realtime_config.validate_environment(headers={}, model=model, api_key=None) async with websockets.connect( url, additional_headers=headers, max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, - ssl=ssl_context, + ssl=vertex_ssl_context, ): return True else: raise ValueError(f"Unsupported model: {model}") - ssl_context = get_shared_realtime_ssl_context() + ssl_context: Final = get_shared_realtime_ssl_context() async with websockets.connect( url, additional_headers=auth_headers, diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index d3d41c5b54b..25eeb0c407f 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch import pytest import litellm +from litellm.models.credentials import CredentialItem from litellm.realtime_api import main as realtime_main from litellm.realtime_api.main import _with_resolved_session_model @@ -224,9 +225,11 @@ def test_client_secret_forwards_nested_transcription_model_untouched(monkeypatch class _CapturingConnect: def __init__(self) -> None: self.url: str | None = None + self.kwargs: dict[str, object] = {} def __call__(self, url: str, **kwargs: object) -> "_CapturingConnect": self.url = url + self.kwargs = kwargs return self async def __aenter__(self) -> MagicMock: @@ -241,6 +244,35 @@ class _CapturingConnect: return None +@pytest.mark.asyncio +async def test_azure_health_check_resolves_stored_credentials(monkeypatch): + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="azure-rt", + credential_values={ + "api_key": "sk-from-credential", + "api_base": "https://example.openai.azure.com", + "api_version": "2025-04-01-preview", + }, + credential_info={}, + ) + ], + ) + connect = _CapturingConnect() + with patch("websockets.connect", connect): + assert await realtime_main._realtime_health_check( + model="gpt-realtime", + custom_llm_provider="azure", + api_key=None, + model_params={"model": "azure/gpt-realtime", "litellm_credential_name": "azure-rt"}, + ) + assert connect.kwargs["additional_headers"] == {"api-key": "sk-from-credential"} + assert connect.url is not None and connect.url.startswith("wss://example.openai.azure.com") + + @pytest.mark.asyncio async def test_azure_health_check_probes_ga_transcription_url_for_transcription_model(local_model_cost_map): """Regression for LIT-6240: transcription-only models (mode audio_transcription From 22b377fe2aad149e0c364e4dab7c17a6440b59ae Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 19:21:03 -0700 Subject: [PATCH 069/107] fix(proxy): log the provider usage on deferred /v1/messages calls and price cache writes without a creation rate With a post-call guardrail the proxy defers async success logging, and every nested wrapper on a /v1/messages call bridged to the Responses API overwrote the stored closure, so the spend log was built from the outermost Anthropic-shaped reply under Responses semantics and recorded the prompt tokens without the cache hit. The first wrapper to exit now keeps the slot, which is the innermost provider response, the same one the non-deferred path logs. The flat cost path also billed cache-creation tokens at 0 when the model had no cache_creation_input_token_cost. It now falls back to the input rate, and the 1h rate to the creation rate, matching the tiered path and the custom pricing helper. --- .../litellm_core_utils/llm_cost_calc/utils.py | 55 ++++--- litellm/utils.py | 68 ++++++--- .../llm_cost_calc/test_llm_cost_calc_utils.py | 55 ++++++- .../test_deferred_guardrail_logging.py | 140 ++++++++++++++++++ 4 files changed, 263 insertions(+), 55 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 8fc428b38ae..3d8edf9e219 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -532,6 +532,11 @@ def _get_token_base_cost( `missing_cache_read_uses_input` resolves an absent cache-read rate to the resolved input rate instead of 0.0; an explicit 0.0 rate stays a real price either way. + An absent cache-creation rate always resolves to the resolved input rate, the way the + tiered table and custom deployment pricing already do, since a provider that publishes + no write price bills cache writes as ordinary input. An absent 1h write rate resolves + to the cache-creation rate. An explicit 0.0 stays a real price for both. + Returns: Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost) """ @@ -554,10 +559,9 @@ def _get_token_base_cost( output_image_cost: Final = _get_cost_per_unit(model_info, "output_cost_per_image_token", None) if output_image_cost is not None: completion_base_cost = cast(float, output_image_cost) - cache_creation_cost = cast(float, _get_cost_per_unit(model_info, cache_creation_cost_key)) - cache_creation_cost_above_1hr = cast( - float, - _get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"), + cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_cost_key, default_value=None) + cache_creation_cost_above_1hr = _get_cost_per_unit( + model_info, "cache_creation_input_token_cost_above_1hr", default_value=None ) cache_read_cost = _get_cost_per_unit(model_info, cache_read_cost_key, default_value=None) @@ -639,22 +643,10 @@ def _get_token_base_cost( else f"cache_read_input_token_cost_above_{threshold_str}_tokens" ) - cache_creation_cost = cast( - float, - _get_cost_per_unit( - model_info, - cache_creation_tiered_key, - cache_creation_cost, - ), - ) + cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_tiered_key, cache_creation_cost) - cache_creation_cost_above_1hr = cast( - float, - _get_cost_per_unit( - model_info, - cache_creation_1hr_tiered_key, - cache_creation_cost_above_1hr, - ), + cache_creation_cost_above_1hr = _get_cost_per_unit( + model_info, cache_creation_1hr_tiered_key, cache_creation_cost_above_1hr ) cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost) @@ -665,16 +657,19 @@ def _get_token_base_cost( except Exception: continue + input_rate_for_missing_cache_rates: Final = _off_peak_rate( + _open_off_peak_block(model_info, current_time) or MappingProxyType({}), + "input_cost_per_token", + prompt_base_cost, + ) if cache_read_cost is None: - cache_read_cost = ( - _off_peak_rate( - _open_off_peak_block(model_info, current_time) or MappingProxyType({}), - "input_cost_per_token", - prompt_base_cost, - ) - if missing_cache_read_uses_input - else 0.0 - ) + cache_read_cost = input_rate_for_missing_cache_rates if missing_cache_read_uses_input else 0.0 + resolved_cache_creation_cost: Final = ( + input_rate_for_missing_cache_rates if cache_creation_cost is None else cache_creation_cost + ) + resolved_cache_creation_cost_above_1hr: Final = ( + resolved_cache_creation_cost if cache_creation_cost_above_1hr is None else cache_creation_cost_above_1hr + ) return _apply_off_peak_to_base_costs( model_info, @@ -682,8 +677,8 @@ def _get_token_base_cost( ( prompt_base_cost, completion_base_cost, - cache_creation_cost, - cache_creation_cost_above_1hr, + resolved_cache_creation_cost, + resolved_cache_creation_cost_above_1hr, cache_read_cost, ), ) diff --git a/litellm/utils.py b/litellm/utils.py index d4e3d58ba9f..62bcea4f468 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1208,30 +1208,13 @@ def _dispatch_success_logging( is_litellm_internal_call: bool, ) -> None: if not is_litellm_internal_call: - if getattr(logging_obj, "_defer_async_logging", False): - - def _enqueue_deferred_logging() -> None: - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) - ) - - logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging - else: - asyncio.create_task( - _client_async_logging_helper( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - ) - ) + _schedule_async_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) logging_obj.handle_sync_success_callbacks_for_async_calls( result=result, @@ -1240,6 +1223,43 @@ def _dispatch_success_logging( ) +def _schedule_async_success_logging( + logging_obj: LiteLLMLoggingObject, + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + is_completion_with_fallbacks: bool, +) -> None: + """Fire the async success log for ``result`` now, or park it on the logging object while + the proxy defers logging past its post-call guardrails. + + Nested @client wrappers (Anthropic Messages over the chat adapter, chat over the Responses + bridge) each exit through here with the same logging object and their own shape of the same + response. The immediate path already logs one request once, since the first task marks + ``has_logged_async_success`` and the later ones skip. The deferred slot keeps the same + first-wins rule: the innermost wrapper's provider-shaped result is the one the spend log + reads usage from, and a later wrapper never swaps in its client-shaped translation. + """ + + def _enqueue_async_logging() -> None: + asyncio.create_task( + _client_async_logging_helper( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + ) + ) + + if not getattr(logging_obj, "_defer_async_logging", False): + _enqueue_async_logging() + return + if getattr(logging_obj, "_enqueue_deferred_logging", None) is not None: + return + logging_obj._enqueue_deferred_logging = _enqueue_async_logging + + async def _client_async_logging_helper( logging_obj: LiteLLMLoggingObject, result, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 2289de9a951..11520b65598 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -4039,7 +4039,7 @@ def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeyp cache_read_input_token_cost=6e-7, cache_read_input_audio_token_cost=6e-7, cache_creation_input_token_cost=7.5e-6, - cache_creation_input_token_cost_above_1hr=0.0, + cache_creation_input_token_cost_above_1hr=7.5e-6, output_cost_per_reasoning_token=3e-5, ) assert breakdown.cache_read_cost == pytest.approx(200_000 * rates.cache_read_input_token_cost) @@ -5334,3 +5334,56 @@ def test_realtime_models_bill_cached_text_and_audio_at_their_cache_read_rates( prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=custom_llm_provider) assert prompt_cost == pytest.approx(expected_prompt_cost) + + +def test_generic_cost_per_token_bills_cache_creation_at_the_input_rate_without_a_write_price(): + """Azure and OpenAI publish no cache-write price and bill cache writes as ordinary input. + A deployment priced with only input, output, and cache-read rates must bill the creation + tokens the provider reports at the input rate, never at 0. The numbers are a cold 7,336-token + prompt on a deployment that reports all but 3 of them as cache creation.""" + model_info = { + "input_cost_per_token": 2e-7, + "output_cost_per_token": 1.25e-6, + "cache_read_input_token_cost": 2e-8, + } + usage = Usage( + prompt_tokens=7336, + completion_tokens=23, + total_tokens=7359, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=0, cache_creation_tokens=7333), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model="custom-priced-deployment", usage=usage, custom_llm_provider="azure", model_info=model_info + ) + + assert prompt_cost == pytest.approx(7336 * 2e-7) + assert completion_cost == pytest.approx(23 * 1.25e-6) + + +@pytest.mark.parametrize( + ("cache_rates", "current_time", "expected_creation", "expected_creation_1h"), + ( + pytest.param({}, None, 2e-7, 2e-7, id="no-write-price-uses-the-input-rate"), + pytest.param({"cache_creation_input_token_cost": 2.5e-7}, None, 2.5e-7, 2.5e-7, id="no-1h-price-uses-the-write-price"), + pytest.param({"cache_creation_input_token_cost": 0.0}, None, 0.0, 0.0, id="explicit-zero-stays-zero"), + pytest.param( + {"off_peak_pricing": {"hours_utc": "00:00-23:59", "input_cost_per_token": 1e-7}}, + datetime(2026, 9, 14, 12, tzinfo=timezone.utc), + 1e-7, + 1e-7, + id="no-write-price-uses-the-off-peak-input-rate", + ), + ), +) +def test_get_token_base_cost_resolves_missing_cache_write_rates_like_the_tiered_path( + cache_rates: dict, current_time: datetime | None, expected_creation: float, expected_creation_1h: float +) -> None: + model_info = {"input_cost_per_token": 2e-7, "output_cost_per_token": 1.25e-6, **cache_rates} + usage = Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11) + + _, _, creation, creation_1h, _ = _get_token_base_cost(model_info, usage, current_time=current_time) + + assert creation == pytest.approx(expected_creation) + assert creation_1h == pytest.approx(expected_creation_1h) + diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index c550a0a41d2..6d254f06ce8 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -16,14 +16,21 @@ Streaming: CSW.__anext__ stores args on logging_obj at stream end. import asyncio import logging +from collections.abc import Callable +from datetime import datetime from typing import Any, Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.types.utils import StandardLoggingPayload +from litellm.utils import _dispatch_success_logging from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -54,6 +61,25 @@ def _attach_mock_success_dispatch(mock_logging_obj, async_success_fn): mock_logging_obj.async_success_handler = async_success_fn +async def _wait_until(condition: Callable[[], bool]) -> None: + """Give the logging worker a bounded window to run what the closure enqueued.""" + for _ in range(200): + if condition(): + return + await asyncio.sleep(0.01) + + +class _RecordingLogger(CustomLogger): + """Keeps what the async success callback was handed, the way a spend logger sees it.""" + + def __init__(self) -> None: + super().__init__() + self.standard_logging_object: StandardLoggingPayload | None = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.standard_logging_object = kwargs["standard_logging_object"] + + class PostCallGuardrail(CustomGuardrail): """A post-call guardrail.""" @@ -259,6 +285,120 @@ async def test_deferred_flag_stores_and_executes_closure(): pass +@pytest.mark.asyncio +async def test_deferred_slot_keeps_the_innermost_wrapper_result(): + """Nested @client wrappers exit through _dispatch_success_logging with one shared logging + object. The deferred slot must keep the first stored result, the way the immediate path's + has_logged dedupe keeps the first fired task, so the spend log reads usage from the + innermost provider-shaped response and never from an outer wrapper's translation of it.""" + logging_obj: Final = MagicMock() + logging_obj._defer_async_logging = True + logging_obj._enqueue_deferred_logging = None + logging_obj.async_success_handler = AsyncMock() + inner_result: Final = object() + outer_result: Final = object() + + for result in (inner_result, outer_result): + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=datetime.now(), + end_time=datetime.now(), + is_completion_with_fallbacks=False, + is_litellm_internal_call=False, + ) + + logging_obj._enqueue_deferred_logging() + await _wait_until(lambda: logging_obj.async_success_handler.await_count > 0) + + logging_obj.async_success_handler.assert_awaited_once() + assert logging_obj.async_success_handler.await_args.kwargs["result"] is inner_result + assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_count == 2 + + +@pytest.mark.asyncio +async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the_provider_usage( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + """/v1/messages on an Azure gpt-5.4+ deployment with function tools runs three nested + wrappers: anthropic_messages, the chat adapter's acompletion, and the Responses bridge + acompletion hands the call to, which retags the call as ``responses``. With logging + deferred for a post-call guardrail the stored closure must carry the innermost provider + response: logging the Anthropic-shaped reply under Responses semantics books this + 7,336-token prompt as 3 tokens, since Anthropic's input_tokens excludes the cache hit.""" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(url__regex=r"https://deferred-nested\.openai\.azure\.com/openai/.*responses.*").mock( + return_value=httpx.Response( + 200, + json={ + "id": "resp_deferred_nested", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-nano", + "output": [ + { + "type": "message", + "id": "msg_deferred_nested", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello!", "annotations": []}], + } + ], + "usage": { + "input_tokens": 7336, + "input_tokens_details": {"cached_tokens": 7333}, + "output_tokens": 23, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 7359, + }, + }, + ) + ) + recorder: Final = _RecordingLogger() + logging_obj: Final = Logging( + model="azure/gpt-5.4-nano", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="deferred-nested-anthropic-messages", + function_id="deferred-nested-anthropic-messages", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj._defer_async_logging = True + + response: Final = await litellm.anthropic_messages( + model="azure/gpt-5.4-nano", + messages=[{"role": "user", "content": "hi"}], + max_tokens=16, + tools=[ + { + "name": "lookup_volume", + "description": "Look up a storage volume by name", + "input_schema": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}, + } + ], + api_key="sk-deferred-nested", + api_base="https://deferred-nested.openai.azure.com", + api_version="2025-04-01-preview", + litellm_logging_obj=logging_obj, + ) + assert response["content"] == [{"type": "text", "text": "Hello!"}] + assert response["usage"]["input_tokens"] == 3 + assert response["usage"]["cache_read_input_tokens"] == 7333 + + logging_obj._enqueue_deferred_logging() + await _wait_until(lambda: recorder.standard_logging_object is not None) + + assert recorder.standard_logging_object is not None + assert recorder.standard_logging_object["prompt_tokens"] == 7336 + assert recorder.standard_logging_object["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 7333 + assert recorder.standard_logging_object["response_cost"] == pytest.approx(3 * 2e-7 + 7333 * 2e-8 + 23 * 1.25e-6) + + # --------------------------------------------------------------------------- # 3. Non-streaming regression: without flag, create_task fires normally # --------------------------------------------------------------------------- From 0a0a9509dcf75c55c138c1d2b6e8f0219f089523 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 19:21:23 -0700 Subject: [PATCH 070/107] test(bedrock): type the session policy test helpers with the policy TypedDict --- .../test_web_identity_session_policy.py | 36 +++++++------------ 1 file changed, 12 insertions(+), 24 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index cd5b14c8d37..2e84c3ca36a 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -37,6 +37,9 @@ from typing import Final from unittest.mock import MagicMock, patch import pytest +from pydantic import TypeAdapter + +from litellm.llms.bedrock.base_aws_llm import WebIdentitySessionPolicy, _SessionPolicyStatement # Actions the Claude Platform on AWS service is documented to call. # Source: AWS IAM action reference + the #27678 surface area. @@ -89,15 +92,18 @@ def _captured_policy_document() -> str: return kwargs["Policy"] -def _captured_policy() -> dict: - return json.loads(_captured_policy_document()) +_SESSION_POLICY_ADAPTER: Final = TypeAdapter(WebIdentitySessionPolicy) -def _granted_actions(policy: dict) -> frozenset[str]: +def _captured_policy() -> WebIdentitySessionPolicy: + return _SESSION_POLICY_ADAPTER.validate_python(json.loads(_captured_policy_document())) + + +def _granted_actions(policy: WebIdentitySessionPolicy) -> frozenset[str]: return frozenset(action for stmt in policy["Statement"] for action in stmt["Action"]) -def _statement_by_sid(policy: dict, sid: str) -> dict: +def _statement_by_sid(policy: WebIdentitySessionPolicy, sid: str) -> _SessionPolicyStatement: for stmt in policy["Statement"]: if stmt.get("Sid") == sid: return stmt @@ -111,7 +117,6 @@ class TestWebIdentitySessionPolicyShape: def test_policy_parses_as_valid_iam_document(self): policy = _captured_policy() assert policy["Version"] == "2012-10-17" - assert isinstance(policy["Statement"], list) assert len(policy["Statement"]) >= 2 def test_bedrock_statement_actions_preserved(self): @@ -146,16 +151,7 @@ class TestClaudePlatformActionsCovered: @pytest.mark.parametrize("action", sorted(_CLAUDE_PLATFORM_ACTIONS)) def test_claude_platform_action_present(self, action: str): - policy = _captured_policy() - # Action may live in any Statement — search across all. - all_actions: set = set() - for stmt in policy["Statement"]: - stmt_actions = stmt.get("Action") - if isinstance(stmt_actions, str): - all_actions.add(stmt_actions) - elif isinstance(stmt_actions, list): - all_actions.update(stmt_actions) - assert action in all_actions, ( + assert action in _granted_actions(_captured_policy()), ( f"{action} missing from session policy — " f"bedrock/claude_platform/* requests will 403 on OIDC auth" ) @@ -188,15 +184,7 @@ class TestBedrockMantleActionsCovered: action" even when the role's identity policy grants it.""" def test_bedrock_mantle_create_inference_present(self): - policy = _captured_policy() - all_actions: set = set() - for stmt in policy["Statement"]: - stmt_actions = stmt.get("Action") - if isinstance(stmt_actions, str): - all_actions.add(stmt_actions) - elif isinstance(stmt_actions, list): - all_actions.update(stmt_actions) - assert "bedrock-mantle:CreateInference" in all_actions, ( + assert "bedrock-mantle:CreateInference" in _granted_actions(_captured_policy()), ( "bedrock-mantle:CreateInference missing from session policy — " "bedrock_mantle/* requests will 403 on OIDC/WIF auth" ) From 123b76e2d80f614a7e3595d74ac7e89d2d7a1335 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:34:49 +0000 Subject: [PATCH 071/107] fix(health): annotate credential resolution casts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/realtime_api/main.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 59aa63eb8a5..336fb3cef12 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -636,9 +636,15 @@ async def _realtime_health_check( import websockets resolved_params: Final = _model_params_with_stored_credentials(model_params or _EMPTY_MODEL_PARAMS) - resolved_api_key: Final = cast(str | None, api_key or resolved_params.get("api_key")) - resolved_api_base: Final = cast(str | None, api_base or resolved_params.get("api_base")) - resolved_api_version: Final = cast(str | None, api_version or resolved_params.get("api_version")) + resolved_api_key: Final = cast( # cast-ok: provider parameters expose optional string credentials + str | None, api_key or resolved_params.get("api_key") + ) + resolved_api_base: Final = cast( # cast-ok: provider parameters expose optional string endpoints + str | None, api_base or resolved_params.get("api_base") + ) + resolved_api_version: Final = cast( # cast-ok: provider parameters expose optional string versions + str | None, api_version or resolved_params.get("api_version") + ) url: str | None = None auth_headers: Final = _realtime_health_check_auth_headers( custom_llm_provider=custom_llm_provider, From e668a313847c9fd946d9a84c7d42e1616e9b9431 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 02:43:48 +0000 Subject: [PATCH 072/107] fix(health): keep realtime credential hydration immutable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/realtime_api/main.py | 11 +++++++---- tests/test_litellm/realtime_api/test_main.py | 5 ++++- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 336fb3cef12..6fdadc52f13 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -28,8 +28,9 @@ from litellm.types.realtime import ( ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, LlmProviders -from litellm.utils import ProviderConfigManager, load_credentials_from_list +from litellm.utils import ProviderConfigManager +from ..litellm_core_utils.credential_accessor import CredentialAccessor from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ..llms.azure.common_utils import get_azure_ad_token @@ -57,9 +58,11 @@ _EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({}) def _model_params_with_stored_credentials(model_params: Mapping[str, Any]) -> Mapping[str, Any]: - hydrated: Final = dict(model_params) - load_credentials_from_list(hydrated) - return MappingProxyType(hydrated) + credential_name: Final = model_params.get("litellm_credential_name") + credential_values: Final = ( + CredentialAccessor.get_credential_values(credential_name) if isinstance(credential_name, str) else {} + ) + return MappingProxyType({**credential_values, **model_params}) def _with_resolved_session_model(session: dict[str, object], model_name: str) -> dict[str, object]: diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 25eeb0c407f..dcd4c01730d 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -267,10 +267,13 @@ async def test_azure_health_check_resolves_stored_credentials(monkeypatch): model="gpt-realtime", custom_llm_provider="azure", api_key=None, + realtime_protocol="beta", model_params={"model": "azure/gpt-realtime", "litellm_credential_name": "azure-rt"}, ) assert connect.kwargs["additional_headers"] == {"api-key": "sk-from-credential"} - assert connect.url is not None and connect.url.startswith("wss://example.openai.azure.com") + assert connect.url is not None + assert connect.url.startswith("wss://example.openai.azure.com") + assert "api-version=2025-04-01-preview" in connect.url @pytest.mark.asyncio From ae90f1a45891debf622d3cb531163e4fb7f21399 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:28:47 -0700 Subject: [PATCH 073/107] fix(guardrails): type the request payload handed to the Anthropic write-back --- .../llms/anthropic/chat/guardrail_translation/handler.py | 7 ++++--- .../test_anthropic_guardrail_handler.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b438d168b52..b8e179e7274 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1079,14 +1079,15 @@ class AnthropicMessagesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - data: dict, - responses: list[str], + data: dict[str, object], # mutable-ok: API message payload + responses: Sequence[str], scanned: tuple[ScannedText, ...], ) -> None: """ Apply guardrail responses back to the top-level system prompt and the input messages. """ - messages: Final[Sequence[_WritableMessage]] = data.get("messages") or () + raw_messages: Final = data.get("messages") + messages: Final[Sequence[_WritableMessage]] = raw_messages if isinstance(raw_messages, list) else () for item, guardrail_response in zip(scanned, responses): match item.target: case SystemStringTarget(): diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 4b47b7d8c4a..51c3751dcf8 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2200,7 +2200,7 @@ class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: inputs, the same way the chat completions handler hands over system messages and tool_calls.""" @staticmethod - def _tool_use_conversation(system): + def _tool_use_conversation(system: str) -> dict[str, Any]: return { "model": "claude-sonnet-4-5", "system": system, From e01d97ea08dd27606672187f029adac680a49b3c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:31:22 -0700 Subject: [PATCH 074/107] fix(guardrails): read Prompt Security modified rows with the slot count's own predicate A chat row whose content carried an empty text part counted two slots in the chat completions handler while Prompt Security read one text out of the modified row, so the structured rewrite was dropped and the request got the named rejection. One shared helper now lists a row's slot texts and both the slot count and the modified-row reader use it. --- .../base_llm/guardrail_translation/utils.py | 12 ++++++---- .../prompt_security/prompt_security.py | 16 ++----------- .../test_prompt_security_guardrails.py | 24 +++++++++++++++++++ 3 files changed, 34 insertions(+), 18 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index c47a56b5ea9..51d43436fc9 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -373,13 +373,17 @@ def _content_part_text(part: object) -> str | None: return text if isinstance(text, str) else None -def message_text_slot_count(message: AllMessageValues) -> int: +def message_slot_texts(message: Mapping[str, object]) -> tuple[str, ...]: content: Final = message.get("content") if isinstance(content, str): - return 1 + return (content,) if isinstance(content, list): - return sum(1 for part in content if _content_part_text(part) is not None) - return 0 + return tuple(text for part in content if (text := _content_part_text(part)) is not None) + return () + + +def message_text_slot_count(message: AllMessageValues) -> int: + return len(message_slot_texts(message)) def _part_with_text(part: object, text: str) -> object: diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 3f29b3e751c..7e43566f224 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -15,7 +15,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) -from litellm.llms.base_llm.guardrail_translation.utils import message_with_slot_texts +from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts, message_with_slot_texts from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -399,19 +399,7 @@ class PromptSecurityGuardrail(CustomGuardrail): return inputs def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]: - """Extract text content from messages.""" - texts: Final = [] - for message in messages: - content = message.get("content") - if isinstance(content, str): - texts.append(content) - elif isinstance(content, list): - for item in content: - if isinstance(item, dict) and item.get("type") == "text": - text = item.get("text") - if text: - texts.append(text) - return texts + return [text for message in messages for text in message_slot_texts(message)] async def _process_standalone_images(self, images: list[str], user_api_key_alias: str | None) -> None: """Process standalone images from inputs (data URLs).""" diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index 785895b6190..3218632a8d2 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -269,6 +269,30 @@ async def test_modify_with_unexpected_message_count_keeps_texts_only(monkeypatch assert result["texts"] == ["Look up [REDACTED]"] +@pytest.mark.asyncio +async def test_modify_keeps_empty_text_parts_as_slots(monkeypatch: pytest.MonkeyPatch): + """The chat handler counts an empty text part as a slot, so a modify verdict + that echoes the empty part still lines up with the row and its texts.""" + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True) + messages: list[AllMessageValues] = [ + {"role": "user", "content": [{"type": "text", "text": "Look up 123-45-6789"}, {"type": "text", "text": ""}]} + ] + inputs = {"texts": ["Look up 123-45-6789", ""], "structured_messages": messages} + modified_messages = [ + {"role": "user", "content": [{"type": "text", "text": "Look up [REDACTED]"}, {"type": "text", "text": ""}]} + ] + + with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)): + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={"messages": messages}, input_type="request" + ) + + assert result["structured_messages"] == modified_messages + assert result["texts"] == ["Look up [REDACTED]", ""] + + @pytest.mark.asyncio async def test_apply_guardrail_allow_request(monkeypatch: pytest.MonkeyPatch): """Test that apply_guardrail allows safe prompts""" From 1111658e16cea78af5509d4dd9a9ca4b0e610b9e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:31:45 -0700 Subject: [PATCH 075/107] fix(bedrock): grant the AgentCore for-user invoke action in the web identity session policy The chat and A2A AgentCore handlers send X-Amzn-Bedrock-AgentCore-Runtime-User-Id when runtimeUserId is set, and AWS requires bedrock-agentcore:InvokeAgentRuntimeForUser alongside InvokeAgentRuntime on that call, so the ceiling now carries both. The role identity policy still decides whether a given role may use it The invalid-token test now uses a neutral example audience --- litellm/llms/bedrock/base_aws_llm.py | 1 + .../llms/bedrock/test_web_identity_session_policy.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 70869f69e3a..385d5898569 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -133,6 +133,7 @@ _WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = Map ), "BedrockAgentCoreLiteLLM": ( "bedrock-agentcore:InvokeAgentRuntime", + "bedrock-agentcore:InvokeAgentRuntimeForUser", "bedrock-agentcore:InvokeGateway", ), "ClaudePlatformLiteLLM": ( diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py index 2e84c3ca36a..cbb69b4ceed 100644 --- a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -230,7 +230,7 @@ class TestInvalidIdentityTokenSurfacesAudience: operator can diagnose the mismatch without enabling LITELLM_LOG=DEBUG on a prod instance.""" - _AUD = "https://guidepoint.litellm-prod.ai" + _AUD = "https://gateway.example.com" _ISS = "https://accounts.google.com" _STS_MESSAGE = ( "An error occurred (InvalidIdentityToken) when calling the " @@ -322,6 +322,9 @@ _BEDROCK_ROUTE_ACTIONS: Final = MappingProxyType( "knowledgebases": "bedrock:ListKnowledgeBases", "agents/{agent_id}/agentAliases/{alias_id}/sessions/{session_id}/text": "bedrock:InvokeAgent", "runtimes/{agent_runtime_arn}/invocations": "bedrock-agentcore:InvokeAgentRuntime", + "runtimes/{agent_runtime_arn}/invocations with X-Amzn-Bedrock-AgentCore-Runtime-User-Id": ( + "bedrock-agentcore:InvokeAgentRuntimeForUser" + ), "mcp": "bedrock-agentcore:InvokeGateway", } ) From bebc76316cd8269c20a5682147745f16bbd5a568 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 14 Sep 2026 21:42:57 -0700 Subject: [PATCH 076/107] fix(cli): label savings cost bars with the auto-router name --- .../proxy/client/cli/commands/statusline_script.py | 5 ++--- .../proxy/client/cli/test_statusline_script.py | 12 +++++++++++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 5493586f627..0e1c1b25e0f 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -50,7 +50,6 @@ CODEX_BASE_URL_ENV_KEYS: Final = ("OPENAI_BASE_URL",) CODEX_API_KEY_ENV_KEYS: Final = ("OPENAI_API_KEY",) CODEX_STOP_EVENT: Final = "Stop" SYNTHETIC_MODEL: Final = "" -LITELLM_LABEL: Final = "LiteLLM" RESET: Final = "\033[0m" BOLD: Final = "\033[1m" DIM: Final = "\033[90m" @@ -316,9 +315,9 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100 delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}") peak: Final = max(session.spend, session.baseline_spend) - label_width: Final = max(len(LITELLM_LABEL), len(reference)) + label_width: Final = max(len(session.router_name), len(reference)) rows: Final = ( - (LITELLM_LABEL, session.spend, LITELLM_COLOR), + (session.router_name, session.spend, LITELLM_COLOR), (reference, session.baseline_spend, BASELINE_COLOR), ) lines: Final = ( diff --git a/tests/test_litellm/proxy/client/cli/test_statusline_script.py b/tests/test_litellm/proxy/client/cli/test_statusline_script.py index 5c0cf6b5703..2b812932542 100644 --- a/tests/test_litellm/proxy/client/cli/test_statusline_script.py +++ b/tests/test_litellm/proxy/client/cli/test_statusline_script.py @@ -253,10 +253,18 @@ class TestRender: text = render("claude-sonnet-5", RECORDED, config_dir, use_color=False, bar_width=10) assert text.splitlines() == [ "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5", - "LiteLLM ████░░░░░░ $0.14", + "claude-auto ████░░░░░░ $0.14", "Claude Opus 5 ██████████ $0.38", ] + def test_a_long_router_name_keeps_both_cost_bars_aligned(self, config_dir: Path) -> None: + session: Final = RECORDED._replace(router_name="engineering-smart-router") + text: Final = render("claude-sonnet-5", session, config_dir, use_color=False, bar_width=10) + assert text.splitlines()[1:] == [ + "engineering-smart-router ████░░░░░░ $0.14", + "Claude Opus 5 ██████████ $0.38", + ] + def test_control_characters_in_any_externally_sourced_label_never_reach_the_terminal(self, tmp_path, config_dir): # The transcript, the proxy payload and Claude Code's model cache all feed labels straight into a # terminal, and none is under this script's control. Only the control bytes are dropped (ESC, BEL, @@ -312,6 +320,7 @@ class TestClaudeCodeMode: text: Final = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) assert text.startswith("claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5\n") + assert text.splitlines()[1].startswith("claude-auto ") def test_a_discovered_display_name_labels_the_sessions_model( self, tmp_path: Path, transcript: Path, config_dir: Path @@ -379,6 +388,7 @@ class TestCodexMode: out = _run({"hook_event_name": "Stop", "session_id": SESSION_ID, "transcript_path": "/nope"}, env, fetch) message = json.loads(out)["systemMessage"] assert message.splitlines()[1] == "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5" + assert message.splitlines()[2].startswith("claude-auto ") assert message.startswith("\n") assert seen == [Credentials("http://127.0.0.1:4000", "sk-codex")] From 31b34f7767cf8426293e98723f8e80cc667a1cab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:46:02 -0700 Subject: [PATCH 077/107] test(guardrails): type the Anthropic write-back test helper --- .../test_anthropic_guardrail_handler.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 51c3751dcf8..32786fc5057 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -13,6 +13,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, @@ -2163,14 +2164,14 @@ class ToolCallArgumentsMaskingGuardrail(InputsRecordingGuardrail): super().__init__() self.return_copies = return_copies self.replacement_arguments = replacement_arguments - self.seen_tool_calls: list[dict] = [] + self.seen_tool_calls: list[dict[str, object]] = [] async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], - logging_obj: Optional[Any] = None, + logging_obj: Optional[LiteLLMLoggingObj] = None, ) -> GenericGuardrailAPIInputs: outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj) tool_calls = list(outputs.get("tool_calls") or []) From f4f1e2eace561ee3b6ed15e916b0e8a20cffd73c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:46:22 -0700 Subject: [PATCH 078/107] refactor(sdk): move the None sentinel to constants and freeze the init kwargs filter --- litellm/constants.py | 2 ++ litellm/exceptions.py | 4 +++- litellm/litellm_core_utils/exception_mapping_utils.py | 7 +++++-- litellm/proxy/common_utils/openai_error_payload.py | 6 +++--- .../litellm_core_utils/test_exception_mapping_utils.py | 10 ---------- .../proxy/common_utils/test_openai_error_payload.py | 3 --- .../proxy/test_common_request_processing.py | 3 --- 7 files changed, 13 insertions(+), 22 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 1211a600747..cbf5efdbca0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1980,6 +1980,8 @@ UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = ( HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS | ORIGIN_SERVER_HEADERS ) +STRINGIFIED_NONE: Final[str] = "None" + # A retrieved response replays the usage of the call that created it, so pricing these # read/management routes like inference bills the same tokens twice. NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset( diff --git a/litellm/exceptions.py b/litellm/exceptions.py index fdc2cc1f169..eb4b5f535ff 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -235,7 +235,9 @@ class BadRequestError(openai.BadRequestError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None + self.headers = ( + {k: str(v) for k, v in headers.items()} if headers else None # mutable-ok: the proxy updates it in place + ) # Use response if it's a valid httpx.Response with a request, otherwise use minimal error response # Note: We check _request (not .request property) to avoid RuntimeError when _request is None if ( diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 36b53a26c99..b3dec655092 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -3,6 +3,7 @@ import json import re import traceback from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final, Protocol, cast import httpx @@ -206,7 +207,7 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None def _accepted_init_kwargs(exception_class: type[Exception], candidates: Mapping[str, object]) -> Mapping[str, object]: accepted: Final = inspect.signature(exception_class).parameters - return {name: value for name, value in candidates.items() if name in accepted} + return MappingProxyType({name: value for name, value in candidates.items() if name in accepted}) def extract_and_raise_litellm_exception( @@ -236,7 +237,9 @@ def extract_and_raise_litellm_exception( message=error_str, llm_provider=custom_llm_provider, model=model, - **_accepted_init_kwargs(raised_exception_obj, {"response": response, "body": body, "headers": headers}), + **_accepted_init_kwargs( + raised_exception_obj, MappingProxyType({"response": response, "body": body, "headers": headers}) + ), ) diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index 90b3c998247..fe23ab2c4b6 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -8,7 +8,7 @@ from typing import Final from fastapi import status -_STRINGIFIED_NONE: Final = "None" +from litellm.constants import STRINGIFIED_NONE _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( { @@ -37,7 +37,7 @@ def openai_error_type(exc: object, status_code: int) -> str: """OpenAI types ``error.type`` as a required string, so an exception carrying none falls back to the type its status code stands for.""" carried: Final = attribute_of(exc, "type") - if isinstance(carried, str) and carried != _STRINGIFIED_NONE: + if isinstance(carried, str) and carried != STRINGIFIED_NONE: return carried mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) if mapped is not None: @@ -51,4 +51,4 @@ def openai_error_param(exc: object) -> str | None: """OpenAI types ``error.param`` as nullable, so an exception carrying none serializes as JSON ``null``.""" carried: Final = attribute_of(exc, "param") - return carried if isinstance(carried, str) and carried != _STRINGIFIED_NONE else None + return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 1ff2bbb9bdd..653c07d06ad 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1430,8 +1430,6 @@ def _openai_handler_error( status_code: int = 400, message: str = _GUARDRAIL_BLOCK_ERROR["message"], ) -> OpenAIError: - """What litellm/llms/openai/openai.py raises after the openai SDK rejects a request: - the SDK's str() carries the wire body, and the handler copies headers and body over.""" wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type, "code": str(status_code), "message": message} return OpenAIError( status_code=status_code, @@ -1448,9 +1446,6 @@ _PROXY_HEADERS = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guar ("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)] ) def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, status_code: int): - """An SDK caller behind a proxy tells a guardrail block from any other 4xx by the body's - provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError - must carry both whichever error.type and status the proxy version on the other end emits.""" with pytest.raises(litellm.BadRequestError) as exc_info: exception_type( model="claude-haiku-4-5", @@ -1469,9 +1464,6 @@ def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, s "relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError] ) def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_class: type[litellm.BadRequestError]): - """A proxy relaying a provider's own litellm error names the class in the message, which - re-raises that class on the SDK side before the generic 400 mapping runs; it must carry the - body and the proxy headers the same way the generic mapping now does.""" message = f"litellm.{relayed_class.__name__}: {_GUARDRAIL_BLOCK_ERROR['message']}" with pytest.raises(relayed_class) as exc_info: @@ -1489,8 +1481,6 @@ def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_clas def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): - """A vendor's own response headers stay on e.response the way every other mapped provider - error keeps them; only a LiteLLM proxy upstream puts headers on e.headers.""" with pytest.raises(litellm.BadRequestError) as exc_info: exception_type( model="gpt-5.4-mini", diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index 90f1da84a61..90850840ab4 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -146,9 +146,6 @@ def test_a_status_carried_by_an_exception_drives_the_type_it_reports(): def test_a_stringified_none_type_or_param_is_treated_as_absent(): - """A proxy fronting a proxy older than 1.102 receives {"type": "None", "param": "None"} on - the wire; the SDK now keeps that body on the mapped exception, and re-emitting the literal - is the exact bug this module exists to stop.""" from litellm.exceptions import BadRequestError carried = BadRequestError( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 1f8fa762fbd..ef5741e472a 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4064,9 +4064,6 @@ class TestHandleLLMApiExceptionFramingHeaders: assert proxy_exc.headers["x-request-id"] == "abc-123" async def test_strips_the_date_and_server_headers_of_an_upstream_litellm_proxy(self): - """A proxy fronting another LiteLLM proxy gets the upstream's date and server - on the mapped exception; forwarding them would duplicate the Date header - uvicorn adds to every response and leak the upstream server identity.""" exc = litellm.BadRequestError( message="Content blocked", llm_provider="litellm_proxy", From 8573241c49817d335de7ec450a7b39deca126356 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:46:33 -0700 Subject: [PATCH 079/107] fix(cost): resolve a missing 1h cache write rate after off-peak pricing The one-hour cache write fallback now takes the applied cache write rate, so an off-peak write price carries into it instead of the input rate The cost estimate test for a cost-map model without cache prices now expects writes at the input rate, which is what the proxy bills The recording logger in the deferred guardrail test types its callback parameters --- .../litellm_core_utils/llm_cost_calc/utils.py | 16 +++++++--------- .../llm_cost_calc/test_llm_cost_calc_utils.py | 13 +++++++++++++ .../test_deferred_guardrail_logging.py | 10 ++++++---- .../test_cost_tracking_settings.py | 15 +++++++++------ 4 files changed, 35 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index fb135f0c60c..baa9aab1087 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -484,11 +484,12 @@ def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None, def _apply_off_peak_to_base_costs( model_info: ModelInfo, current_time: datetime | None, - base_costs: tuple[float, float, float, float, float], + base_costs: tuple[float, float, float, float | None, float], ) -> tuple[float, float, float, float, float]: """Apply off-peak rates to an already-resolved set of base costs, whichever pricing path - produced them. The one-hour cache-creation rate passes through untouched, since - off_peak_pricing has no field for it, and reasoning is left to _resolve_billed_reasoning_rate. + produced them. off_peak_pricing has no field for the one-hour cache-creation rate, so a + present one passes through untouched and an absent one resolves to the applied + cache-creation rate. Reasoning is left to _resolve_billed_reasoning_rate. """ prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs rates: Final = apply_off_peak_pricing( @@ -506,7 +507,7 @@ def _apply_off_peak_to_base_costs( rates.input_rate, rates.output_rate, rates.cache_creation_rate, - cache_creation_above_1hr, + rates.cache_creation_rate if cache_creation_above_1hr is None else cache_creation_above_1hr, rates.cache_read_rate, ) @@ -535,7 +536,7 @@ def _get_token_base_cost( An absent cache-creation rate always resolves to the resolved input rate, the way the tiered table and custom deployment pricing already do, since a provider that publishes no write price bills cache writes as ordinary input. An absent 1h write rate resolves - to the cache-creation rate. An explicit 0.0 stays a real price for both. + to the cache-creation rate, off-peak included. An explicit 0.0 stays a real price for both. Returns: Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost) @@ -667,9 +668,6 @@ def _get_token_base_cost( resolved_cache_creation_cost: Final = ( input_rate_for_missing_cache_rates if cache_creation_cost is None else cache_creation_cost ) - resolved_cache_creation_cost_above_1hr: Final = ( - resolved_cache_creation_cost if cache_creation_cost_above_1hr is None else cache_creation_cost_above_1hr - ) return _apply_off_peak_to_base_costs( model_info, @@ -678,7 +676,7 @@ def _get_token_base_cost( prompt_base_cost, completion_base_cost, resolved_cache_creation_cost, - resolved_cache_creation_cost_above_1hr, + cache_creation_cost_above_1hr, cache_read_cost, ), ) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 508402f514d..fb406a8a7a6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -5476,6 +5476,19 @@ def test_generic_cost_per_token_bills_cache_creation_at_the_input_rate_without_a 1e-7, id="no-write-price-uses-the-off-peak-input-rate", ), + pytest.param( + { + "off_peak_pricing": { + "hours_utc": "00:00-23:59", + "input_cost_per_token": 1e-7, + "cache_creation_input_token_cost": 3e-7, + } + }, + datetime(2026, 9, 14, 12, tzinfo=timezone.utc), + 3e-7, + 3e-7, + id="no-1h-price-uses-the-off-peak-write-price", + ), ), ) def test_get_token_base_cost_resolves_missing_cache_write_rates_like_the_tiered_path( diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 6d254f06ce8..6295469c066 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -16,9 +16,9 @@ Streaming: CSW.__anext__ stores args on logging_obj at stream end. import asyncio import logging -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime -from typing import Any, Final +from typing import Any, Final, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -76,8 +76,10 @@ class _RecordingLogger(CustomLogger): super().__init__() self.standard_logging_object: StandardLoggingPayload | None = None - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.standard_logging_object = kwargs["standard_logging_object"] + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.standard_logging_object = cast(StandardLoggingPayload, kwargs["standard_logging_object"]) class PostCallGuardrail(CustomGuardrail): diff --git a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py index 7ece35ceedf..c73d29e78b2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cost_tracking_settings.py @@ -975,8 +975,9 @@ class TestEstimateCostCacheAndReasoningTokens: @pytest.mark.asyncio async def test_a_model_without_cache_or_reasoning_prices_estimates_what_the_proxy_bills(self, monkeypatch): - """The cost calculator bills cache tokens of a cost-map model without cache prices at zero - and its reasoning tokens at the output rate. The estimate reports those effective rates.""" + """The cost calculator bills cache reads of a cost-map model without cache prices at zero, + its cache writes at the input rate, and its reasoning tokens at the output rate. The estimate + reports those effective rates.""" monkeypatch.setitem( litellm.model_cost, A_MAPPED_MODEL, @@ -986,12 +987,14 @@ class TestEstimateCostCacheAndReasoningTokens: response = await _estimate_with_cache_and_reasoning(None, model=A_MAPPED_MODEL) assert response.cache_read_cost_per_request == 0.0 - assert response.cache_creation_cost_per_request == 0.0 + assert response.cache_creation_cost_per_request == pytest.approx(CACHE_CREATION_TOKENS * 5e-6) assert response.reasoning_cost_per_request == pytest.approx(REASONING_TOKENS * 6e-6) - assert response.input_cost_per_request == pytest.approx(TEXT_INPUT_TOKENS * 5e-6) - assert response.cost_per_request == pytest.approx(TEXT_INPUT_TOKENS * 5e-6 + OUTPUT_TOKENS * 6e-6) + assert response.input_cost_per_request == pytest.approx((TEXT_INPUT_TOKENS + CACHE_CREATION_TOKENS) * 5e-6) + assert response.cost_per_request == pytest.approx( + (TEXT_INPUT_TOKENS + CACHE_CREATION_TOKENS) * 5e-6 + OUTPUT_TOKENS * 6e-6 + ) assert response.cache_read_input_token_cost == 0.0 - assert response.cache_creation_input_token_cost == 0.0 + assert response.cache_creation_input_token_cost == pytest.approx(5e-6) assert response.output_cost_per_reasoning_token == pytest.approx(6e-6) @pytest.mark.asyncio From 7e3d7178b4194e642fe4a9a3a5ffae79945f8612 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 21:52:01 -0700 Subject: [PATCH 080/107] test(spend): reconcile concurrent requests and daily activity --- .../spend_tracking/spend_reconciliation.py | 109 +++++++++++++ .../spend_tracking/test_spend_tracking_e2e.py | 56 ++----- .../test_team_daily_activity_e2e.py | 144 +++++++++++++++--- 3 files changed, 244 insertions(+), 65 deletions(-) create mode 100644 tests/e2e/quota_management/spend_tracking/spend_reconciliation.py diff --git a/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py new file mode 100644 index 00000000000..8fcfea3f296 --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py @@ -0,0 +1,109 @@ +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from math import isclose +from typing import Final + +from e2e_config import provider_edge_base, unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody, LiteLLMParamsBody, TeamNewBody +from spend_e2e_client import SpendClient + +INPUT_RATE: Final = 0.00004 +OUTPUT_RATE: Final = 0.00008 + + +@dataclass(frozen=True) +class TeamTraffic: + team_id: str + key: str + responses: tuple[ChatResponse, ...] + + @property + def prompt_tokens(self) -> int: + return sum(response.usage.prompt_tokens or 0 for response in self.responses if response.usage) + + @property + def completion_tokens(self) -> int: + return sum(response.usage.completion_tokens or 0 for response in self.responses if response.usage) + + @property + def spend(self) -> float: + return self.prompt_tokens * INPUT_RATE + self.completion_tokens * OUTPUT_RATE + + +def create_traffic(client: SpendClient, resources: ResourceManager) -> tuple[TeamTraffic, ...]: + base: Final = provider_edge_base("openai") + model: Final = f"e2e-reconciliation-{unique_marker()}" + model_id: Final = client.proxy.create_model( + model, + LiteLLMParamsBody( + model="openai/gpt-5.6-luna", + api_key="os.environ/OPENAI_API_KEY", + api_base=None if base is None else f"{base}/v1", + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + def team_traffic() -> TeamTraffic: + team: Final = client.proxy.create_team(TeamNewBody(team_alias=f"e2e-spend-{unique_marker()}")) + resources.defer(lambda: client.proxy.delete_team(team)) + key: Final = client.proxy.generate_key(KeyGenerateBody(team_id=team, models=[model])) + resources.defer(lambda: client.proxy.delete_key(key)) + + prompts: Final = tuple(f"Reply with one word. {index} {unique_marker()}" for index in range(7)) + + def call(index: int) -> ChatResponse: + response: Final = unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=prompts[index])], + max_completion_tokens=128, + ), + ) + ) + assert response.id, "successful response must have an ID" + assert response.usage is not None, "successful response must have usage" + assert response.usage.prompt_tokens is not None and response.usage.prompt_tokens > 0 + assert response.usage.completion_tokens is not None and response.usage.completion_tokens > 0 + assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + assert not response.usage.cache_creation_input_tokens + assert not response.usage.cache_read_input_tokens + assert not response.usage.prompt_tokens_details or not response.usage.prompt_tokens_details.cached_tokens + return response + + sequential: Final = call(0) + with ThreadPoolExecutor(max_workers=6) as pool: + concurrent: Final = tuple(pool.map(call, range(1, 7))) + return TeamTraffic(team, key, (sequential, *concurrent)) + + return tuple(team_traffic() for _ in range(2)) + + +def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None: + expected_ids: Final = frozenset(response.id for response in traffic.responses) + assert len(expected_ids) == len(traffic.responses), "responses must have distinct IDs" + rows: Final = client.poll_logs_for_key( + traffic.key, + min_rows=len(traffic.responses), + predicate=lambda values: frozenset(row.request_id for row in values) == expected_ids, + ) + assert frozenset(row.request_id for row in rows) == expected_ids, "stored IDs must equal returned response IDs" + assert len(rows) == len(traffic.responses), "expected exactly one scoped spend row per response" + by_id: Final = {row.request_id: row for row in rows} + for response in traffic.responses: + row = by_id[response.id] + usage = response.usage + assert usage is not None and usage.prompt_tokens is not None and usage.completion_tokens is not None + assert row.team_id == traffic.team_id + assert row.status == "success" + assert row.cache_hit != "True" + assert row.prompt_tokens == usage.prompt_tokens + assert row.completion_tokens == usage.completion_tokens + assert row.total_tokens == usage.total_tokens + expected_cost = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE + assert row.spend is not None and isclose(row.spend, expected_cost, rel_tol=1e-6, abs_tol=1e-9) diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py index c5d76d44580..750285cbfb6 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py @@ -17,13 +17,12 @@ fails the test; a pricing or token-count drift does not. import time from collections.abc import Callable -from concurrent.futures import ThreadPoolExecutor +from math import isclose import pytest - -from e2e_http import Result, Success +from e2e_http import Success from lifecycle import ResourceManager -from models import ChatResponse, LiteLLMParamsBody, SpendLogs, SpendLogsParams +from models import LiteLLMParamsBody, SpendLogs, SpendLogsParams from spend_e2e_client import SpendClient, SpendLogRow, is_ok, unique_marker, unwrap pytestmark = pytest.mark.e2e @@ -280,51 +279,18 @@ def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> N ), f"key aggregate {key_spend} != sum of logs {logs_total}; rows: {_summarize(rows)}" +@pytest.mark.replayable @pytest.mark.covers("quota_management.spend_tracking.concurrent_burst.loses_no_spend") def test_burst_of_concurrent_calls_loses_no_spend( - client: SpendClient, scoped_key: str + client: SpendClient, resources: ResourceManager ) -> None: - """Six concurrent calls on one key: every call lands its own spend row under a - distinct request_id and the key aggregate equals the sum of the rows. - Sequential accuracy is covered by test_key_spend_equals_sum_of_logs; this pins - the concurrent increment path (parallel writers racing on one key's counter), - where a lost update can never be reproduced by sequential calls.""" - burst = 6 + from spend_reconciliation import assert_logs_match, create_traffic - def call(idx: int) -> Result[ChatResponse]: - return client.chat( - scoped_key, - "gemini-2.5-flash", - f"burst call {idx} {unique_marker()}", - max_tokens=16, - ) - - with ThreadPoolExecutor(max_workers=burst) as pool: - results = tuple(pool.map(call, range(burst))) - failed = [r for r in results if not is_ok(r)] - assert not failed, f"{len(failed)}/{burst} burst calls failed; first: {failed[0]}" - - rows = client.poll_logs_for_key( - scoped_key, - min_rows=burst, - predicate=lambda rs: len([r for r in rs if (r.spend or 0) > 0]) >= burst, - ) - costed = [r for r in rows if (r.spend or 0) > 0] - assert len(costed) >= burst, ( - f"only {len(costed)}/{burst} burst calls produced a costed row - " - f"rows lost under concurrency: {_summarize(rows)}" - ) - request_ids = [r.request_id for r in costed] - assert len(set(request_ids)) == len(request_ids), ( - f"concurrent rows collapsed onto shared request_ids: {_summarize(rows)}" - ) - - logs_total = sum((r.spend or 0) for r in rows) - key_spend = client.poll_key_spend(scoped_key, minimum=logs_total * 0.999) - assert _approx_equal(key_spend, logs_total), ( - f"key aggregate {key_spend} != sum of {len(rows)} rows {logs_total} - " - f"spend increments lost under concurrency: {_summarize(rows)}" - ) + traffic = create_traffic(client, resources) + for team in traffic: + assert_logs_match(client, team) + key_spend = client.poll_key_spend(team.key, minimum=team.spend * 0.999999) + assert isclose(key_spend, team.spend, rel_tol=1e-6, abs_tol=1e-9) @pytest.mark.covers("quota_management.spend_tracking.pagination.keeps_total") diff --git a/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py b/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py index ed0a6af4ec9..dca5572f510 100644 --- a/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py @@ -7,13 +7,17 @@ missing start/end dates are rejected. from __future__ import annotations +import time from datetime import datetime, timedelta, timezone +from math import isclose +from typing import Final import pytest from e2e_http import ProbeResult -from models import DateRangeParams +from lifecycle import ResourceManager from pydantic import BaseModel from spend_e2e_client import SpendClient +from spend_reconciliation import assert_logs_match, create_traffic pytestmark = pytest.mark.e2e @@ -24,22 +28,45 @@ class TeamDailyActivityParams(BaseModel): start_date: str | None = None end_date: str | None = None page: int = 1 + page_size: int = 1 + team_ids: str | None = None class TeamDailyActivityRow(BaseModel): date: str metrics: TeamDailyActivityMetrics + breakdown: TeamDailyActivityBreakdown class TeamDailyActivityMetrics(BaseModel): spend: float total_tokens: int + prompt_tokens: int + completion_tokens: int + api_requests: int + successful_requests: int + failed_requests: int + + +class TeamDailyActivityEntity(BaseModel): + metrics: TeamDailyActivityMetrics + + +class TeamDailyActivityBreakdown(BaseModel): + entities: dict[str, TeamDailyActivityEntity] class TeamDailyActivityMetadata(BaseModel): page: int total_pages: int has_more: bool + total_spend: float + total_prompt_tokens: int + total_completion_tokens: int + total_tokens: int + total_api_requests: int + total_successful_requests: int + total_failed_requests: int class TeamDailyActivityResponse(BaseModel): @@ -47,32 +74,109 @@ class TeamDailyActivityResponse(BaseModel): metadata: TeamDailyActivityMetadata -def _range_days(days: int) -> DateRangeParams: - end = datetime.now(timezone.utc).date() - start = end - timedelta(days=days) - return DateRangeParams(start_date=start.isoformat(), end_date=end.isoformat()) - - def _probe(client: SpendClient, params: BaseModel) -> ProbeResult: return client.proxy.transport.probe(ROUTE, params=params) class TestTeamDailyActivity: + @pytest.mark.replayable @pytest.mark.covers("mgmt.team.daily_activity.happy_path") - @pytest.mark.parametrize("days", [1, 7, 30]) - def test_valid_date_range_returns_results_and_metadata(self, client: SpendClient, days: int) -> None: - result = _probe(client, _range_days(days)) - assert result.status_code == 200, ( - f"{ROUTE} range={days}d must be 200, got {result.status_code}: {result.body[:600]}" + def test_valid_date_range_returns_results_and_metadata( + self, client: SpendClient, resources: ResourceManager + ) -> None: + started: Final = datetime.now(timezone.utc).date() + traffic: Final = create_traffic(client, resources) + for team in traffic: + assert_logs_match(client, team) + ended: Final = datetime.now(timezone.utc).date() + team_ids: Final = ",".join(team.team_id for team in traffic) + + def fetch( + page: int, start: str = started.isoformat(), end: str = ended.isoformat() + ) -> TeamDailyActivityResponse: + result: Final = _probe( + client, + TeamDailyActivityParams( + start_date=start, + end_date=end, + page=page, + page_size=1, + team_ids=team_ids, + ), + ) + assert result.status_code == 200, f"daily activity failed: {result.status_code} {result.body[:300]}" + return TeamDailyActivityResponse.model_validate_json(result.body) + + def pages() -> tuple[TeamDailyActivityResponse, ...]: + first: Final = fetch(1) + assert first.metadata.total_pages <= len(traffic) * 2, "unexpected extra scoped daily groups" + return (first, *(fetch(page) for page in range(2, first.metadata.total_pages + 1))) + + deadline: Final = time.monotonic() + client.proxy.poll_timeout + while True: + observed = pages() + if sum(page.metadata.total_api_requests for page in observed) >= sum(len(t.responses) for t in traffic): + break + if time.monotonic() >= deadline: + break + time.sleep(client.proxy.poll_interval) + + assert len(observed) >= 2, "two teams must exercise a page boundary" + for index, page in enumerate(observed, 1): + assert page.metadata.page == index + assert page.metadata.total_pages == len(observed) + assert page.metadata.has_more == (index < len(observed)) + assert len(page.results) == 1, "each fetched daily group must appear in results" + row = page.results[0] + assert started <= datetime.fromisoformat(row.date).date() <= ended + assert len(row.breakdown.entities) == 1 + assert row.metrics.total_tokens == page.metadata.total_tokens + assert row.metrics.prompt_tokens == page.metadata.total_prompt_tokens + assert row.metrics.completion_tokens == page.metadata.total_completion_tokens + assert row.metrics.api_requests == page.metadata.total_api_requests + assert row.metrics.successful_requests == page.metadata.total_successful_requests + assert row.metrics.failed_requests == page.metadata.total_failed_requests + assert isclose(row.metrics.spend, page.metadata.total_spend, rel_tol=1e-6, abs_tol=1e-9) + + entities: Final = tuple( + (team_id, entity.metrics) + for page in observed + for row in page.results + for team_id, entity in row.breakdown.entities.items() ) - parsed = TeamDailyActivityResponse.model_validate_json(result.body) - assert parsed.metadata.page == 1 - assert parsed.metadata.total_pages >= 1 - if parsed.results: - first = parsed.results[0] - assert first.date - assert first.metrics.spend >= 0 - assert first.metrics.total_tokens >= 0 + assert frozenset(team_id for team_id, _ in entities) == frozenset(team.team_id for team in traffic) + for team in traffic: + metrics = tuple(metrics for team_id, metrics in entities if team_id == team.team_id) + assert sum(m.api_requests for m in metrics) == len(team.responses) + assert sum(m.successful_requests for m in metrics) == len(team.responses) + assert sum(m.failed_requests for m in metrics) == 0 + assert sum(m.prompt_tokens for m in metrics) == team.prompt_tokens + assert sum(m.completion_tokens for m in metrics) == team.completion_tokens + assert sum(m.total_tokens for m in metrics) == team.prompt_tokens + team.completion_tokens + assert isclose(sum(m.spend for m in metrics), team.spend, rel_tol=1e-6, abs_tol=1e-9) + assert isclose( + sum(page.metadata.total_spend for page in observed), + sum(team.spend for team in traffic), + rel_tol=1e-6, + abs_tol=1e-9, + ) + assert sum(page.metadata.total_tokens for page in observed) == sum( + team.prompt_tokens + team.completion_tokens for team in traffic + ) + + empty_date: Final = (started - timedelta(days=7)).isoformat() + empty: Final = fetch(1, empty_date, empty_date) + assert empty.results == [] + assert empty.metadata.total_pages == 0 + assert empty.metadata.page == 1 + assert not empty.metadata.has_more + assert empty.metadata.total_spend == 0 + assert empty.metadata.total_tokens == 0 + assert empty.metadata.total_api_requests == 0 + assert empty.metadata.total_prompt_tokens == 0 + assert empty.metadata.total_completion_tokens == 0 + assert empty.metadata.total_successful_requests == 0 + assert empty.metadata.total_failed_requests == 0 @pytest.mark.covers("mgmt.team.daily_activity.missing_start_date_rejected") def test_missing_start_date_is_rejected(self, client: SpendClient) -> None: From d7a2bdd4415898630145037c077a61b22d932913 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:56:23 -0700 Subject: [PATCH 081/107] test(cost): type the cache rate cases of the base cost test --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index fb406a8a7a6..30b158e3b2c 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2,6 +2,7 @@ import json from datetime import datetime, timezone import pytest +from collections.abc import Mapping from fastapi.testclient import TestClient import litellm @@ -5492,7 +5493,10 @@ def test_generic_cost_per_token_bills_cache_creation_at_the_input_rate_without_a ), ) def test_get_token_base_cost_resolves_missing_cache_write_rates_like_the_tiered_path( - cache_rates: dict, current_time: datetime | None, expected_creation: float, expected_creation_1h: float + cache_rates: Mapping[str, float | Mapping[str, float | str]], + current_time: datetime | None, + expected_creation: float, + expected_creation_1h: float, ) -> None: model_info = {"input_cost_per_token": 2e-7, "output_cost_per_token": 1.25e-6, **cache_rates} usage = Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11) From 714b113c5fa43e1c6ee147656513200e526219eb Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 15 Sep 2026 05:10:32 +0000 Subject: [PATCH 082/107] fix(guardrails): leave scoped-out image-only input unrecorded and split the not_run helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/guardrail_translation/handler.py | 53 +++++++++++-------- .../test_openai_guardrail_handler.py | 20 +++++++ 2 files changed, 52 insertions(+), 21 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 332285722e1..20ca6a95aab 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -214,27 +214,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) elif not images_to_check and not guardrail_to_apply.records_own_guardrail_information: - unscoped_texts: Final[list[str]] = [] - unscoped_tool_calls: Final[list[ChatCompletionToolParam]] = [] - for unscoped_idx, unscoped_message in enumerate(messages): - self._extract_inputs( - message=unscoped_message, - msg_idx=unscoped_idx, - texts_to_check=unscoped_texts, - images_to_check=[], - tool_calls_to_check=unscoped_tool_calls, - text_task_mappings=[], - tool_call_task_mappings=[], - ) - guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response=( - "no scannable content after message scoping" - if unscoped_texts or unscoped_tool_calls - else "no scannable content" - ), - request_data=data, - guardrail_status="not_run", - ) + self._record_not_run(data=data, messages=messages, guardrail_to_apply=guardrail_to_apply) verbose_proxy_logger.debug( "OpenAI Chat Completions: Processed input messages: %s", @@ -243,6 +223,37 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return data + def _record_not_run( + self, + data: dict, + messages: list[dict[str, Any]], + guardrail_to_apply: "CustomGuardrail", + ) -> None: + unscoped_texts: Final[list[str]] = [] + unscoped_images: Final[list[str]] = [] + unscoped_tool_calls: Final[list[ChatCompletionToolParam]] = [] + for msg_idx, message in enumerate(messages): + self._extract_inputs( + message=message, + msg_idx=msg_idx, + texts_to_check=unscoped_texts, + images_to_check=unscoped_images, + tool_calls_to_check=unscoped_tool_calls, + text_task_mappings=[], + tool_call_task_mappings=[], + ) + if unscoped_images: + return + guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=( + "no scannable content after message scoping" + if unscoped_texts or unscoped_tool_calls + else "no scannable content" + ), + request_data=data, + guardrail_status="not_run", + ) + def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name).""" names: Final[list[str]] = [] diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 617fba587b2..b176d1e9057 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1980,6 +1980,26 @@ class TestNoScannableContentRecordsNotRun: assert self._recorded_entries(data) == [] + @pytest.mark.asyncio + async def test_scoped_out_image_only_message_is_not_reported_as_not_run(self): + """An image in a skipped role must behave like any other image-only request""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="image-guardrail") + guardrail.skip_system_message_in_guardrail = True + data = { + "messages": [ + { + "role": "system", + "content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}], + }, + ] + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.last_inputs is None + assert self._recorded_entries(data) == [] + class ToolDroppingTextGuardrail(CustomGuardrail): """Answers one text per non-tool message it saw, the way a guardrail that From fe0fb97fd2d0a39ad022739547f464b48a7a061b Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 15 Sep 2026 05:20:56 +0000 Subject: [PATCH 083/107] fix(guardrails): record not_run when a skipped role mixes text and images Only image-only unscoped content stays unrecorded; text or tool content removed by scoping is recorded as not_run even when an image sits beside it. Also keeps the type-discipline budget flat by returning the reason from the helper and annotating the accumulator lists _extract_inputs requires. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat/guardrail_translation/handler.py | 51 +++++++++---------- litellm/proxy/compliance_checks.py | 2 +- .../test_openai_guardrail_handler.py | 26 ++++++++++ 3 files changed, 52 insertions(+), 27 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 20ca6a95aab..01e14f2248d 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -213,8 +213,16 @@ class OpenAIChatCompletionsHandler(BaseTranslation): task_mappings=tool_call_task_mappings, ) - elif not images_to_check and not guardrail_to_apply.records_own_guardrail_information: - self._record_not_run(data=data, messages=messages, guardrail_to_apply=guardrail_to_apply) + elif ( + not images_to_check + and not guardrail_to_apply.records_own_guardrail_information + and (not_run_reason := self._not_run_reason(messages)) is not None + ): + guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=not_run_reason, + request_data=data, + guardrail_status="not_run", + ) verbose_proxy_logger.debug( "OpenAI Chat Completions: Processed input messages: %s", @@ -223,36 +231,27 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return data - def _record_not_run( + def _not_run_reason( self, - data: dict, - messages: list[dict[str, Any]], - guardrail_to_apply: "CustomGuardrail", - ) -> None: - unscoped_texts: Final[list[str]] = [] - unscoped_images: Final[list[str]] = [] - unscoped_tool_calls: Final[list[ChatCompletionToolParam]] = [] + messages: Sequence[dict[str, Any]], # mutable-ok: raw request messages consumed by _extract_inputs + ) -> str | None: + """Why nothing was scanned, or None when the only unscoped content is images, which this handler never scans.""" + texts: Final[list[str]] = [] # mutable-ok: filled by _extract_inputs + images: Final[list[str]] = [] # mutable-ok: filled by _extract_inputs + tool_calls: Final[list[ChatCompletionToolParam]] = [] # mutable-ok: filled by _extract_inputs for msg_idx, message in enumerate(messages): self._extract_inputs( message=message, msg_idx=msg_idx, - texts_to_check=unscoped_texts, - images_to_check=unscoped_images, - tool_calls_to_check=unscoped_tool_calls, - text_task_mappings=[], - tool_call_task_mappings=[], + texts_to_check=texts, + images_to_check=images, + tool_calls_to_check=tool_calls, + text_task_mappings=[], # mutable-ok: required by _extract_inputs, unused here + tool_call_task_mappings=[], # mutable-ok: required by _extract_inputs, unused here ) - if unscoped_images: - return - guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response=( - "no scannable content after message scoping" - if unscoped_texts or unscoped_tool_calls - else "no scannable content" - ), - request_data=data, - guardrail_status="not_run", - ) + if texts or tool_calls: + return "no scannable content after message scoping" + return None if images else "no scannable content" def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name).""" diff --git a/litellm/proxy/compliance_checks.py b/litellm/proxy/compliance_checks.py index d9cc1d0f4fc..9d2f2dc7c69 100644 --- a/litellm/proxy/compliance_checks.py +++ b/litellm/proxy/compliance_checks.py @@ -26,7 +26,7 @@ class ComplianceChecker: def __init__(self, data: ComplianceCheckRequest): self.data = data - self.guardrails = [g for g in (data.guardrail_information or []) if g.get("guardrail_status") != "not_run"] + self.guardrails = tuple(g for g in data.guardrail_information or () if g.get("guardrail_status") != "not_run") def _get_guardrails_by_mode(self, mode: str) -> list[dict]: """ diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index b176d1e9057..cb884fb7cc1 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -2000,6 +2000,32 @@ class TestNoScannableContentRecordsNotRun: assert guardrail.last_inputs is None assert self._recorded_entries(data) == [] + @pytest.mark.asyncio + async def test_scoped_out_text_with_image_records_not_run(self): + """Scoping removed text too, so the skip is recorded even though an image sat beside it""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="image-guardrail") + guardrail.skip_system_message_in_guardrail = True + data = { + "messages": [ + { + "role": "system", + "content": [ + {"type": "text", "text": "Describe this picture."}, + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}, + ], + }, + ] + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.last_inputs is None + entries = self._recorded_entries(data) + assert len(entries) == 1 + assert entries[0]["guardrail_status"] == "not_run" + assert entries[0]["guardrail_response"] == "no scannable content after message scoping" + class ToolDroppingTextGuardrail(CustomGuardrail): """Answers one text per non-tool message it saw, the way a guardrail that From 6a635cbb64fdf7f567bb26058321447b92fb859e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:22:26 -0700 Subject: [PATCH 084/107] fix(sdk): carry a litellm_proxy error's headers on e.response, not e.headers A mapped litellm_proxy exception now attaches an httpx.Response that carries the proxy's response headers whenever the handler attached a header-less synthetic one, on every status branch and on the relay path. BadRequestError keeps its base-class contract: .headers stays the proxy-supplied channel, so the proxy edge keeps forwarding an upstream proxy's headers under the llm_provider- prefix and the date and server edge change is no longer needed. --- litellm/constants.py | 6 +- litellm/exceptions.py | 8 +- .../exception_mapping_utils.py | 64 ++- .../test_exception_mapping_utils.py | 121 ++--- .../proxy/test_common_request_processing.py | 455 ++++++------------ 5 files changed, 223 insertions(+), 431 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index cbf5efdbca0..09442d6151e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1974,11 +1974,7 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( } ) -ORIGIN_SERVER_HEADERS: Final[frozenset[str]] = frozenset({"date", "server"}) - -UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = ( - HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS | ORIGIN_SERVER_HEADERS -) +UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS STRINGIFIED_NONE: Final[str] = "None" diff --git a/litellm/exceptions.py b/litellm/exceptions.py index eb4b5f535ff..3f22a4b2dcd 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -10,7 +10,7 @@ ## LiteLLM versions of the OpenAI Exception Types import enum -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from typing import Any, Final import httpx @@ -226,7 +226,6 @@ class BadRequestError(openai.BadRequestError): max_retries: int | None = None, num_retries: int | None = None, body: dict | None = None, - headers: Mapping[str, str] | None = None, ): self.status_code = 400 self.message = f"litellm.BadRequestError: {message}" @@ -235,9 +234,6 @@ class BadRequestError(openai.BadRequestError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - self.headers = ( - {k: str(v) for k, v in headers.items()} if headers else None # mutable-ok: the proxy updates it in place - ) # Use response if it's a valid httpx.Response with a request, otherwise use minimal error response # Note: We check _request (not .request property) to avoid RuntimeError when _request is None if ( @@ -628,7 +624,6 @@ class ContentPolicyViolationError(BadRequestError): litellm_debug_info: str | None = None, provider_specific_fields: dict | None = None, body: dict | None = None, - headers: Mapping[str, str] | None = None, ): self.status_code = 400 self.message = f"litellm.ContentPolicyViolationError: {message}" @@ -643,7 +638,6 @@ class ContentPolicyViolationError(BadRequestError): response=response, litellm_debug_info=self.litellm_debug_info, body=body, - headers=headers, ) # Call the base class constructor with the parameters it needs def __str__(self): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index b3dec655092..e8406b87777 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -216,7 +216,6 @@ def extract_and_raise_litellm_exception( model: str, custom_llm_provider: str, body: object | None = None, - headers: Mapping[str, str] | None = None, ): """ Covers scenario where litellm sdk calling proxy. @@ -237,9 +236,7 @@ def extract_and_raise_litellm_exception( message=error_str, llm_provider=custom_llm_provider, model=model, - **_accepted_init_kwargs( - raised_exception_obj, MappingProxyType({"response": response, "body": body, "headers": headers}) - ), + **_accepted_init_kwargs(raised_exception_obj, MappingProxyType({"response": response, "body": body})), ) @@ -253,13 +250,20 @@ class _ProviderHTTPException(Protocol): llm_provider: str -def _litellm_proxy_response_headers( +def _litellm_proxy_response( original_exception: _ProviderHTTPException, custom_llm_provider: str -) -> Mapping[str, str] | None: - if custom_llm_provider != "litellm_proxy": - return None +) -> httpx.Response | None: + response: Final = getattr(original_exception, "response", None) + if custom_llm_provider != "litellm_proxy" or not isinstance(response, httpx.Response) or response.headers: + return response headers: Final = getattr(original_exception, "headers", None) - return headers if isinstance(headers, Mapping) else None + if not isinstance(headers, Mapping) or not headers: + return response + return httpx.Response( + status_code=response.status_code, + headers={str(k): str(v) for k, v in headers.items()}, + request=getattr(original_exception, "request", None), + ) def _map_openai_exception( @@ -272,7 +276,7 @@ def _map_openai_exception( exception_provider: str, extra_information: str, ) -> None: - upstream_headers: Final = _litellm_proxy_response_headers(original_exception, custom_llm_provider) + response: Final = _litellm_proxy_response(original_exception, custom_llm_provider) # custom_llm_provider is openai, make it OpenAI message = get_error_message(error_obj=original_exception) if message is None: @@ -301,14 +305,14 @@ def _map_openai_exception( message=f"RateLimitError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, ) elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): raise ContextWindowExceededError( message=f"ContextWindowExceededError: {exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif "invalid_request_error" in error_str and "model_not_found" in error_str: @@ -316,7 +320,7 @@ def _map_openai_exception( message=f"{exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif "A timeout occurred" in error_str: @@ -335,10 +339,9 @@ def _map_openai_exception( message=f"ContentPolicyViolationError: {exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), - headers=upstream_headers, ) elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str: helpful_message: Final = ( @@ -356,20 +359,18 @@ def _map_openai_exception( message=helpful_message, llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), - headers=upstream_headers, ) elif "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str: raise BadRequestError( message=f"{exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), - headers=upstream_headers, ) elif ( "Web server is returning an unknown error" in error_str @@ -385,7 +386,7 @@ def _map_openai_exception( message=f"RateLimitError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif ( @@ -396,7 +397,7 @@ def _map_openai_exception( message=f"AuthenticationError: {exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif "Mistral API raised a streaming error" in error_str: @@ -415,17 +416,16 @@ def _map_openai_exception( message=f"{exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), - headers=upstream_headers, ) elif original_exception.status_code == 401: raise AuthenticationError( message=f"AuthenticationError: {exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 404: @@ -433,7 +433,7 @@ def _map_openai_exception( message=f"NotFoundError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 408: @@ -448,17 +448,16 @@ def _map_openai_exception( message=f"{exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), - headers=upstream_headers, ) elif original_exception.status_code == 429: raise RateLimitError( message=f"RateLimitError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 500: @@ -466,7 +465,7 @@ def _map_openai_exception( message=f"InternalServerError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 502: @@ -474,7 +473,7 @@ def _map_openai_exception( message=f"BadGatewayError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 503: @@ -482,7 +481,7 @@ def _map_openai_exception( message=f"ServiceUnavailableError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 504: # gateway timeout error @@ -2439,12 +2438,11 @@ def exception_type( custom_llm_provider == "litellm_proxy" ): # handle special case where calling litellm proxy + exception str contains error message extract_and_raise_litellm_exception( - response=getattr(original_exception, "response", None), + response=_litellm_proxy_response(mappable_exception, custom_llm_provider), error_str=error_str, model=model, custom_llm_provider=custom_llm_provider, body=getattr(original_exception, "body", None), - headers=_litellm_proxy_response_headers(mappable_exception, custom_llm_provider), ) if ( custom_llm_provider == "openai" diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 653c07d06ad..a4869aac8fe 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1,4 +1,3 @@ - import httpx import openai import pytest @@ -178,9 +177,7 @@ class TestExceptionCheckers: ] for error_str in error_strings: - result = ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) + result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) assert result is True, f"Should detect policy violation in: {error_str}" def test_is_azure_content_policy_violation_error_case_insensitive(self): @@ -194,12 +191,8 @@ class TestExceptionCheckers: ] for error_str in error_strings: - result = ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) - assert ( - result is True - ), f"Should detect policy violation in uppercase: {error_str}" + result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) + assert result is True, f"Should detect policy violation in uppercase: {error_str}" def test_is_azure_content_policy_violation_error_with_non_policy_errors(self): """Test that non-policy violation errors are not detected as policy violations""" @@ -216,12 +209,8 @@ class TestExceptionCheckers: ] for error_str in error_strings: - result = ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) - assert ( - result is False - ), f"Should NOT detect policy violation in: {error_str}" + result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) + assert result is False, f"Should NOT detect policy violation in: {error_str}" def test_is_azure_content_policy_violation_error_with_partial_matches(self): """Test that partial keyword matches work correctly""" @@ -234,9 +223,7 @@ class TestExceptionCheckers: ] for error_str in positive_cases: - result = ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) + result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) assert result is True, f"Should detect policy violation in: {error_str}" # These should not match even though they contain similar words @@ -248,12 +235,8 @@ class TestExceptionCheckers: ] for error_str in negative_cases: - result = ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) - assert ( - result is False - ), f"Should NOT detect policy violation in: {error_str}" + result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) + assert result is False, f"Should NOT detect policy violation in: {error_str}" gemini_context_window_test_cases = [ @@ -271,12 +254,8 @@ gemini_context_window_test_cases = [ ] -@pytest.mark.parametrize( - "error_message, should_raise_context_window", gemini_context_window_test_cases -) -def test_gemini_context_window_error_mapping( - error_message, should_raise_context_window -): +@pytest.mark.parametrize("error_message, should_raise_context_window", gemini_context_window_test_cases) +def test_gemini_context_window_error_mapping(error_message, should_raise_context_window): """ Tests that the exception_type function correctly maps Gemini's context window exceeded errors to litellm.ContextWindowExceededError. @@ -421,9 +400,7 @@ vertex_rate_limit_test_cases = [ ] -@pytest.mark.parametrize( - "error_message, should_raise_rate_limit", vertex_rate_limit_test_cases -) +@pytest.mark.parametrize("error_message, should_raise_rate_limit", vertex_rate_limit_test_cases) def test_vertex_ai_rate_limit_error_mapping(error_message, should_raise_rate_limit): """ Tests that the exception_type function correctly maps Vertex AI's @@ -458,10 +435,7 @@ class TestGetBodyErrorCode: """Unit tests for _get_body_error_code helper.""" def test_parses_int_code(self): - body = ( - '{"error":{"message":"high demand","type":"upstream_error",' - '"param":"","code":429}}' - ) + body = '{"error":{"message":"high demand","type":"upstream_error","param":"","code":429}}' assert _get_body_error_code(body) == 429 def test_parses_string_code(self): @@ -498,8 +472,7 @@ gemini_body_code_429_test_cases = [ ), ( 503, - '{"error":{"message":"upstream unavailable","type":"upstream_error",' - '"param":"","code":429}}', + '{"error":{"message":"upstream unavailable","type":"upstream_error","param":"","code":429}}', litellm.RateLimitError, "HTTP 503 envelope with body code:429 -> RateLimitError", ), @@ -769,9 +742,7 @@ class _UpstreamHTTPError(Exception): self.message = "upstream failure" self.status_code = status_code self.request = httpx.Request("POST", "https://api.example.com/v1/chat/completions") - self.response = httpx.Response( - status_code=status_code, request=self.request, text="upstream failure" - ) + self.response = httpx.Response(status_code=status_code, request=self.request, text="upstream failure") UPSTREAM_STATUS_CODES = (400, 401, 403, 404, 408, 422, 429, 500, 503) @@ -892,15 +863,13 @@ PROVIDERS_WITHOUT_A_HANDLER = tuple( MINIMAX_401_BODY = ( '{"type":"error","error":{"type":"authorized_error","message":"login fail: Please carry the API secret key ' - "in the 'Authorization' field of the request header (1004)\",\"http_code\":\"401\"}," + 'in the \'Authorization\' field of the request header (1004)","http_code":"401"},' '"request_id":"06ddc9ba97ee6340e38f10e09787f547"}' ) def _expected_for(provider: str, status_code: int) -> tuple[type[Exception], int]: - return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get( - status_code, OPENAI_SHAPED[status_code] - ) + return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get(status_code, OPENAI_SHAPED[status_code]) @pytest.fixture @@ -910,9 +879,7 @@ def quiet_exception_mapping(monkeypatch): @pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES) @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_an_upstream_status_maps_to_one_exception_per_provider( - provider, status_code, quiet_exception_mapping -): +def test_an_upstream_status_maps_to_one_exception_per_provider(provider, status_code, quiet_exception_mapping): expected_class, expected_status = _expected_for(provider, status_code) with pytest.raises(openai.APIError) as raised: @@ -928,9 +895,7 @@ def test_an_upstream_status_maps_to_one_exception_per_provider( @pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES) @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from( - provider, status_code, quiet_exception_mapping -): +def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(provider, status_code, quiet_exception_mapping): with pytest.raises(openai.APIError) as raised: exception_type( model="test-model", @@ -943,12 +908,8 @@ def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from( @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_an_already_mapped_litellm_exception_passes_through_untouched( - provider, quiet_exception_mapping -): - already_mapped = litellm.RateLimitError( - message="already mapped", llm_provider=provider, model="test-model" - ) +def test_an_already_mapped_litellm_exception_passes_through_untouched(provider, quiet_exception_mapping): + already_mapped = litellm.RateLimitError(message="already mapped", llm_provider=provider, model="test-model") returned = exception_type( model="test-model", @@ -961,9 +922,7 @@ def test_an_already_mapped_litellm_exception_passes_through_untouched( @pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES) @pytest.mark.parametrize("provider", PROVIDERS_WITHOUT_A_HANDLER) -def test_a_provider_without_a_handler_maps_by_the_upstream_status( - provider, status_code, quiet_exception_mapping -): +def test_a_provider_without_a_handler_maps_by_the_upstream_status(provider, status_code, quiet_exception_mapping): expected_class, expected_status = STATUS_KEYED[status_code] with pytest.raises(openai.APIError) as raised: @@ -1015,9 +974,7 @@ def test_an_unmapped_exception_with_no_model_or_provider_is_a_connection_error(q assert "boom" in raised.value.message -def _raise_and_map( - model: str | None, original_exception: Exception, custom_llm_provider: str | None -) -> None: +def _raise_and_map(model: str | None, original_exception: Exception, custom_llm_provider: str | None) -> None: """Calls exception_type() from inside the except block, as litellm/main.py does, so traceback.format_exc() has a real stack.""" try: @@ -1058,9 +1015,7 @@ def test_an_unmapped_exception_with_no_model_or_provider_message_keeps_traceback CONTEXT_WINDOW_MESSAGE = "This model's maximum context length is 4096 tokens." -CONTENT_POLICY_MESSAGE = ( - '{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}' -) +CONTENT_POLICY_MESSAGE = '{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}' TIMEOUT_MESSAGE = "Request timed out." PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW = ( @@ -1103,15 +1058,11 @@ class _UpstreamErrorWithMessage(_UpstreamHTTPError): super().__init__(status_code=status_code) self.args = (message,) self.message = message - self.response = httpx.Response( - status_code=status_code, request=self.request, text=message - ) + self.response = httpx.Response(status_code=status_code, request=self.request, text=message) @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it( - provider, quiet_exception_mapping -): +def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(provider, quiet_exception_mapping): if provider in PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW: expected_class, expected_status = litellm.ContextWindowExceededError, 400 else: @@ -1129,9 +1080,7 @@ def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it( @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it( - provider, quiet_exception_mapping -): +def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(provider, quiet_exception_mapping): if provider in PROVIDERS_THAT_RECOGNISE_A_CONTENT_POLICY_BLOCK: expected_class, expected_status = litellm.ContentPolicyViolationError, 400 else: @@ -1149,9 +1098,7 @@ def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it( @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_a_timed_out_request_is_a_timeout_for_every_provider( - provider, quiet_exception_mapping -): +def test_a_timed_out_request_is_a_timeout_for_every_provider(provider, quiet_exception_mapping): with pytest.raises(litellm.Timeout) as raised: exception_type( model="test-model", @@ -1442,9 +1389,7 @@ def _openai_handler_error( _PROXY_HEADERS = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} -@pytest.mark.parametrize( - ("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)] -) +@pytest.mark.parametrize(("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)]) def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, status_code: int): with pytest.raises(litellm.BadRequestError) as exc_info: exception_type( @@ -1457,12 +1402,10 @@ def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, s assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" assert exc_info.value.body["type"] == error_type - assert exc_info.value.headers == _PROXY_HEADERS + assert dict(exc_info.value.response.headers) == _PROXY_HEADERS -@pytest.mark.parametrize( - "relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError] -) +@pytest.mark.parametrize("relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError]) def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_class: type[litellm.BadRequestError]): message = f"litellm.{relayed_class.__name__}: {_GUARDRAIL_BLOCK_ERROR['message']}" @@ -1477,7 +1420,7 @@ def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_clas assert type(exc_info.value) is relayed_class assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" - assert exc_info.value.headers == _PROXY_HEADERS + assert dict(exc_info.value.response.headers) == _PROXY_HEADERS def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): @@ -1491,4 +1434,4 @@ def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): ) assert exc_info.value.body["type"] == "vendor_specific_error" - assert exc_info.value.headers is None + assert not exc_info.value.response.headers diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index ef5741e472a..a33b491fce2 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -127,16 +127,12 @@ class TestProxyBaseLLMRequestProcessing: assert json.loads(result.body) == guardrailed_body @pytest.mark.asyncio - async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers( - self, monkeypatch - ): + async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch): """The guardrail JSON path must forward upstream response headers (e.g. x-amzn-requestid) alongside the x-litellm-* headers, matching the non-guardrail passthrough path, while dropping length headers that no longer match the rewritten body.""" - processing_obj = ProxyBaseLLMRequestProcessing( - data={"custom_llm_provider": "bedrock"} - ) + processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -176,14 +172,10 @@ class TestProxyBaseLLMRequestProcessing: assert result.headers["content-length"] == str(len(result.body)) @pytest.mark.asyncio - async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers( - self, monkeypatch - ): + async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch): """The guardrail event-stream branch must also forward upstream response headers alongside the x-litellm-* headers.""" - processing_obj = ProxyBaseLLMRequestProcessing( - data={"custom_llm_provider": "bedrock"} - ) + processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -225,15 +217,11 @@ class TestProxyBaseLLMRequestProcessing: assert result.headers["x-litellm-call-id"] == "test-call-id" @pytest.mark.asyncio - async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook( - self, monkeypatch - ): + async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(self, monkeypatch): """Guardrailed non-streaming passthrough responses must include headers injected by post_call_response_headers_hook, matching the headers a non-guardrailed passthrough response would carry.""" - processing_obj = ProxyBaseLLMRequestProcessing( - data={"custom_llm_provider": "bedrock"} - ) + processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -252,9 +240,7 @@ class TestProxyBaseLLMRequestProcessing: return kwargs["response"] proxy_logging_obj.post_call_success_hook = fake_post_call_success_hook - proxy_logging_obj.post_call_response_headers_hook = AsyncMock( - return_value={"x-litellm-custom": "from-hook"} - ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-litellm-custom": "from-hook"}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=upstream, @@ -378,9 +364,7 @@ class TestProxyBaseLLMRequestProcessing: json.dumps(persisted_body) @pytest.mark.asyncio - async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails( - self, monkeypatch - ): + async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails(self, monkeypatch): """arm_pre_call must run before pre_call_hook: an auto router's own compression policy has to be in `data["metadata"]` (naming the model-side guardrail so it runs even if it isn't default_on) by the time guardrails see the request.""" @@ -2191,16 +2175,10 @@ class TestCommonRequestProcessingHelpers: def _stringified_none_paths(node: object, path: str = "error") -> tuple[str, ...]: if isinstance(node, dict): - return tuple( - found - for key, value in node.items() - for found in _stringified_none_paths(value, f"{path}.{key}") - ) + return tuple(found for key, value in node.items() for found in _stringified_none_paths(value, f"{path}.{key}")) if isinstance(node, (list, tuple)): return tuple( - found - for index, value in enumerate(node) - for found in _stringified_none_paths(value, f"{path}[{index}]") + found for index, value in enumerate(node) for found in _stringified_none_paths(value, f"{path}[{index}]") ) return (path,) if node == "None" else () @@ -2947,9 +2925,7 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={}, - litellm_logging_obj=self._timing_logging_obj( - {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} - ), + litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), ) assert headers["x-litellm-response-duration-ms"] == "500.0" @@ -2968,9 +2944,7 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={}, - litellm_logging_obj=self._timing_logging_obj( - {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} - ), + litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), read_timing_from_logging_obj=False, ) @@ -2991,9 +2965,7 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={"_response_ms": 300.0}, - litellm_logging_obj=self._timing_logging_obj( - {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} - ), + litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), ) assert headers["x-litellm-response-duration-ms"] == "300.0" @@ -3034,9 +3006,7 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={"_response_ms": 300.0, "litellm_overhead_time_ms": 7.5}, - litellm_logging_obj=self._timing_logging_obj( - {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} - ), + litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), ) assert headers["x-litellm-response-duration-ms"] == "300.0" @@ -3488,9 +3458,7 @@ class TestStreamCloseOnDisconnect: finally: closed.set() - response = _UpstreamClosingStreamingResponse( - body(), media_type="text/event-stream" - ) + response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream") async def receive(): await asyncio.Event().wait() @@ -3521,9 +3489,7 @@ class TestStreamCloseOnDisconnect: finally: closed.set() - response = _UpstreamClosingStreamingResponse( - body(), media_type="text/event-stream" - ) + response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream") async def receive(): await disconnected.wait() @@ -3594,9 +3560,7 @@ class TestStreamCloseOnDisconnect: finally: inner_closed.set() - response = await create_response( - generator=wrapped(), media_type="text/event-stream", headers={} - ) + response = await create_response(generator=wrapped(), media_type="text/event-stream", headers={}) async def receive(): await asyncio.Event().wait() @@ -3828,9 +3792,7 @@ class TestStreamCloseOnDisconnect: with pytest.raises(_ClientDisconnectedBeforeFirstChunk): await asyncio.wait_for( - _buffer_first_chunk_honoring_disconnect( - AcloseRaises(), request=self._request_that_disconnects() - ), + _buffer_first_chunk_honoring_disconnect(AcloseRaises(), request=self._request_that_disconnects()), timeout=5, ) @@ -3846,9 +3808,7 @@ class TestStreamCloseOnDisconnect: with pytest.raises(_ClientDisconnectedBeforeFirstChunk): await asyncio.wait_for( - _buffer_first_chunk_honoring_disconnect( - blocking_gen(), request=self._request_that_disconnects() - ), + _buffer_first_chunk_honoring_disconnect(blocking_gen(), request=self._request_that_disconnects()), timeout=5, ) assert closed.is_set() @@ -3864,9 +3824,7 @@ class TestHandleLLMApiExceptionRetryAfter: user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock( - return_value=callback_headers or {} - ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=callback_headers or {}) try: await processor._handle_llm_api_exception( @@ -3918,9 +3876,7 @@ class TestHandleLLMApiExceptionRetryAfter: enable_pre_call_checks=False, cooldown_list=[], ) - proxy_exc = await self._invoke( - exc, callback_headers={"retry-after": "", "x-custom": "1"} - ) + proxy_exc = await self._invoke(exc, callback_headers={"retry-after": "", "x-custom": "1"}) assert proxy_exc.headers["retry-after"] == "43" assert proxy_exc.headers["x-custom"] == "1" @@ -4063,18 +4019,6 @@ class TestHandleLLMApiExceptionFramingHeaders: assert proxy_exc.headers["x-custom-safe"] == "1" assert proxy_exc.headers["x-request-id"] == "abc-123" - async def test_strips_the_date_and_server_headers_of_an_upstream_litellm_proxy(self): - exc = litellm.BadRequestError( - message="Content blocked", - llm_provider="litellm_proxy", - model="claude-haiku-4-5", - headers={"date": "Sun, 13 Sep 2026 08:43:51 GMT", "server": "uvicorn", "x-request-id": "abc-123"}, - ) - proxy_exc = await self._invoke(exc) - assert "date" not in proxy_exc.headers - assert "server" not in proxy_exc.headers - assert proxy_exc.headers["x-request-id"] == "abc-123" - class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" @@ -4161,9 +4105,7 @@ class TestDisconnectGatherCleanup: return Request(scope={"type": "http", "headers": []}, receive=receive) @pytest.mark.asyncio - async def test_base_process_llm_request_raises_499_on_client_disconnect( - self, monkeypatch - ): + async def test_base_process_llm_request_raises_499_on_client_disconnect(self, monkeypatch): """With cancel_on_disconnect enabled, base_process_llm_request returns 499.""" import asyncio @@ -4192,9 +4134,7 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr( - processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) - ) + monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) with pytest.raises(HTTPException) as exc_info: await processing_obj.base_process_llm_request( @@ -4212,9 +4152,7 @@ class TestDisconnectGatherCleanup: assert "disconnected" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect( - self, monkeypatch - ): + async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(self, monkeypatch): import asyncio import litellm.proxy.common_request_processing as cpr @@ -4239,9 +4177,7 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr( - processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) - ) + monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) monkeypatch.setattr( cpr, "route_request", @@ -4302,9 +4238,7 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr( - processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) - ) + monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) with pytest.raises(HTTPException): await processing_obj.base_process_llm_request( @@ -4355,9 +4289,7 @@ class TestDisconnectGatherCleanup: assert task.done() @pytest.mark.asyncio - async def test_base_process_llm_request_preserves_llm_error_after_gather( - self, monkeypatch - ): + async def test_base_process_llm_request_preserves_llm_error_after_gather(self, monkeypatch): import litellm.proxy.common_request_processing as cpr from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -4386,9 +4318,7 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr( - processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) - ) + monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) mock_request = MagicMock(spec=Request) mock_request.is_disconnected = AsyncMock(return_value=False) @@ -4425,19 +4355,13 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": {}}, } - recorded = await _record_streaming_client_disconnect_if_needed( - mock_request, request_data - ) + recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True + assert request_data["metadata"]["error_information"]["error_code"] == "499" assert ( - request_data["metadata"]["error_information"]["error_code"] == "499" - ) - assert ( - mock_logging_obj.model_call_details["litellm_params"]["metadata"][ - "error_information" - ]["error_code"] + mock_logging_obj.model_call_details["litellm_params"]["metadata"]["error_information"]["error_code"] == "499" ) @@ -4451,9 +4375,7 @@ class TestStreamingClientDisconnectLogging: mock_request.is_disconnected = AsyncMock(return_value=False) request_data = {"metadata": {}} - recorded = await _record_streaming_client_disconnect_if_needed( - mock_request, request_data - ) + recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) assert recorded is False assert "client_disconnected" not in request_data["metadata"] @@ -4478,22 +4400,12 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": {}}, } - recorded = await _record_streaming_client_disconnect_if_needed( - mock_request, request_data - ) + recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert ( - mock_logging_obj.model_call_details["litellm_params"]["metadata"][ - "client_disconnected" - ] - is True - ) - assert ( - mock_logging_obj.model_call_details["metadata"]["client_disconnected"] - is True - ) + assert mock_logging_obj.model_call_details["litellm_params"]["metadata"]["client_disconnected"] is True + assert mock_logging_obj.model_call_details["metadata"]["client_disconnected"] is True @pytest.mark.asyncio async def test_record_streaming_client_disconnect_handles_none_request_data_metadata(self): @@ -4509,15 +4421,11 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": None}, } - recorded = await _record_streaming_client_disconnect_if_needed( - mock_request, request_data - ) + recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert ( - request_data["litellm_params"]["metadata"]["client_disconnected"] is True - ) + assert request_data["litellm_params"]["metadata"]["client_disconnected"] is True @pytest.mark.asyncio async def test_apply_client_disconnect_metadata_none_returns_early(self): @@ -4528,9 +4436,7 @@ class TestStreamingClientDisconnectLogging: _apply_client_disconnect_metadata(None) @pytest.mark.asyncio - async def test_finalize_streaming_generator_cleanup_fires_deferred_logging( - self, monkeypatch - ): + async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(self, monkeypatch): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4562,9 +4468,7 @@ class TestStreamingClientDisconnectLogging: assert request_data["metadata"]["error_information"]["error_code"] == "499" @pytest.mark.asyncio - async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion( - self, monkeypatch - ): + async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(self, monkeypatch): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4594,9 +4498,7 @@ class TestStreamingClientDisconnectLogging: assert "client_disconnected" not in request_data["metadata"] @pytest.mark.asyncio - async def test_async_streaming_data_generator_records_499_on_early_aclose( - self, monkeypatch - ): + async def test_async_streaming_data_generator_records_499_on_early_aclose(self, monkeypatch): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4611,9 +4513,7 @@ class TestStreamingClientDisconnectLogging: yield {"choices": [{"delta": {"content": " there"}}]} mock_proxy_logging = MagicMock(spec=ProxyLogging) - mock_proxy_logging.async_post_call_streaming_iterator_hook = ( - mock_streaming_iterator - ) + mock_proxy_logging.async_post_call_streaming_iterator_hook = mock_streaming_iterator ProxyLogging._callback_capabilities_cache.clear() mock_request = MagicMock(spec=Request) @@ -4624,9 +4524,7 @@ class TestStreamingClientDisconnectLogging: "model": "gemini-2.0-flash", "metadata": {}, "litellm_params": {"metadata": {}}, - "litellm_logging_obj": MagicMock( - model_call_details={"metadata": {}, "litellm_params": {}} - ), + "litellm_logging_obj": MagicMock(model_call_details={"metadata": {}, "litellm_params": {}}), } gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( @@ -4645,6 +4543,8 @@ class TestStreamingClientDisconnectLogging: assert request_data["metadata"]["error_information"]["error_code"] == "499" ProxyLogging._callback_capabilities_cache.clear() + + class TestCancelOnDisconnect: """ Coverage for the opt-in `general_settings.cancel_on_disconnect` flag: @@ -4671,23 +4571,17 @@ class TestCancelOnDisconnect: llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - await _cancel_llm_call_on_client_disconnect( - request, llm_call, disconnect_event - ) + await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) assert llm_call.cancelled() assert disconnect_event.is_set() async def test_monitor_is_noop_while_client_stays_connected(self): - request = self._request( - [{"type": "http.request", "body": b"", "more_body": False}] - ) + request = self._request([{"type": "http.request", "body": b"", "more_body": False}]) llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - monitor = asyncio.create_task( - _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) - ) + monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)) await asyncio.sleep(0.01) assert not monitor.done() @@ -4706,9 +4600,7 @@ class TestCancelOnDisconnect: llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - await _cancel_llm_call_on_client_disconnect( - request, llm_call, disconnect_event - ) + await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) assert not llm_call.cancelled() assert not disconnect_event.is_set() @@ -4723,9 +4615,7 @@ class TestCancelOnDisconnect: with pytest.raises(asyncio.CancelledError): await _await_llm_call_cancelling_on_disconnect(request, llm_call) - async def _drive_base_process_llm_request( - self, monkeypatch, general_settings: dict, llm_call, request: Request - ): + async def _drive_base_process_llm_request(self, monkeypatch, general_settings: dict, llm_call, request: Request): from litellm.proxy._types import UserAPIKeyAuth logging_obj = MagicMock() @@ -4734,9 +4624,7 @@ class TestCancelOnDisconnect: logging_obj._on_deferred_stream_complete = None logging_obj.cost_breakdown = None - processor = ProxyBaseLLMRequestProcessing( - data={"model": "fake-model", "litellm_logging_obj": logging_obj} - ) + processor = ProxyBaseLLMRequestProcessing(data={"model": "fake-model", "litellm_logging_obj": logging_obj}) proxy_logging_obj = MagicMock(spec=ProxyLogging) proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) @@ -4744,9 +4632,7 @@ class TestCancelOnDisconnect: proxy_logging_obj.post_call_success_hook = AsyncMock( side_effect=lambda data, user_api_key_dict, response: response ) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock( - return_value=None - ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None) async def fake_route_request(**kwargs): return llm_call() @@ -4825,9 +4711,7 @@ class TestCancelOnDisconnect: with pytest.raises(ProxyException) as exc_info: await processor._handle_llm_api_exception( - e=HTTPException( - status_code=499, detail="Client disconnected the request" - ), + e=HTTPException(status_code=499, detail="Client disconnected the request"), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), proxy_logging_obj=proxy_logging_obj, ) @@ -4893,7 +4777,9 @@ class TestAllmPassthroughRoutePostCallGuardrails: proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", capture_hook) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True + ): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4943,7 +4829,9 @@ class TestAllmPassthroughRoutePostCallGuardrails: proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", non_dict_hook) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True + ): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4981,7 +4869,9 @@ class TestAllmPassthroughRoutePostCallGuardrails: hook_spy = AsyncMock() monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True + ): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -5022,7 +4912,9 @@ class TestAllmPassthroughRoutePostCallGuardrails: hook_spy = AsyncMock() monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False + ): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -5134,7 +5026,9 @@ class TestEventStreamAllmPassthroughRoute: "content-length": "99", } - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True + ): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=mock_response, @@ -5165,9 +5059,7 @@ class TestAllmPassthroughStreamingProviderGate: de-anonymized. """ - def _build_processing_obj( - self, custom_llm_provider: str, endpoint: str = "" - ) -> ProxyBaseLLMRequestProcessing: + def _build_processing_obj(self, custom_llm_provider: str, endpoint: str = "") -> ProxyBaseLLMRequestProcessing: logging_obj = MagicMock() logging_obj.litellm_call_id = "call-123" logging_obj.cost_breakdown = None @@ -5254,14 +5146,17 @@ class TestAllmPassthroughStreamingProviderGate: processing_obj = self._build_processing_obj("anthropic") chunks = [b"chunk-1", b"chunk-2"] - with patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, + with ( + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5270,27 +5165,27 @@ class TestAllmPassthroughStreamingProviderGate: assert streamed == chunks @pytest.mark.asyncio - async def test_bedrock_converse_stream_is_buffered_through_handler( - self, monkeypatch - ): - processing_obj = self._build_processing_obj( - "bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream" - ) + async def test_bedrock_converse_stream_is_buffered_through_handler(self, monkeypatch): + processing_obj = self._build_processing_obj("bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream") chunks = [b"raw-1", b"raw-2"] - with patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=b"modified-body"), - ) as mock_handler: + with ( + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), + patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler, + ): result = await self._run(processing_obj, monkeypatch, chunks) assert isinstance(result, Response) @@ -5306,19 +5201,23 @@ class TestAllmPassthroughStreamingProviderGate: ) chunks = [b"raw-1", b"raw-2"] - with patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=b"modified-body"), - ) as mock_handler: + with ( + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), + patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler, + ): result = await self._run(processing_obj, monkeypatch, chunks) assert isinstance(result, StreamingResponse) @@ -5340,14 +5239,17 @@ class TestAllmPassthroughStreamingProviderGate: ) chunks = [b"raw-1", b"raw-2"] - with patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=False, + with ( + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, + ), ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5366,14 +5268,17 @@ class TestAllmPassthroughStreamingProviderGate: processing_obj = self._build_processing_obj("anthropic") chunks = [b"chunk-1", b"chunk-2"] - with patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=False, + with ( + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, + ), ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5821,9 +5726,7 @@ class TestCostHeadersForCallsPricedAtZero: fastapi_response = Response() processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj}) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False): await processing_obj.base_process_llm_request( request=MagicMock(spec=Request, headers={}), fastapi_response=fastapi_response, @@ -5894,9 +5797,7 @@ class TestCostHeadersForCallsPricedAtZero: assert breakdown.tool_usage_cost == 0.0 def test_cost_breakdown_stays_empty_for_an_inference_call(self): - breakdown = _get_cost_breakdown_from_logging_obj( - litellm_logging_obj=self._logging_obj(call_type="acompletion") - ) + breakdown = _get_cost_breakdown_from_logging_obj(litellm_logging_obj=self._logging_obj(call_type="acompletion")) assert breakdown == CostBreakdownHeaderValues() @@ -5923,7 +5824,6 @@ class TestCostHeadersForCallsPricedAtZero: class TestPreCallWithFallbacksOnLocalRateLimit: - @pytest.mark.asyncio async def test_fallback_triggered_on_local_rate_limit(self): from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -6075,9 +5975,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] user_api_key_dict = MagicMock() - user_api_key_dict.router_settings = { - "fallbacks": [{"gpt-4": ["claude-3-haiku"]}] - } + user_api_key_dict.router_settings = {"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]} with patch.object( processor, @@ -6108,9 +6006,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - processor = ProxyBaseLLMRequestProcessing( - data={"model": "gpt-4", "disable_fallbacks": True} - ) + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "disable_fallbacks": True}) async def mock_pre_call_logic(**kwargs): raise ProxyRateLimitError( @@ -6236,9 +6132,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Real per-key per-model TPM limiter + a key carrying the customer's # `model_tpm_limit` metadata (only the primary is capped). - limiter = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) user_api_key_dict = UserAPIKeyAuth( api_key="sk-lit3890", metadata={"model_tpm_limit": {primary_model: 100}}, @@ -6246,10 +6140,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Pre-seed the primary's per-model token counter at the cap so the very # next request trips it. The counter key uses the *hashed* api_key. - counter_key = ( - f"{user_api_key_dict.api_key}::{primary_model}" - f"::{precise_minute}::request_count" - ) + counter_key = f"{user_api_key_dict.api_key}::{primary_model}::{precise_minute}::request_count" await limiter.internal_usage_cache.async_set_cache( key=counter_key, value={"current_requests": 0, "current_tpm": 100, "current_rpm": 0}, @@ -6280,9 +6171,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router = MagicMock() mock_router.fallbacks = [{primary_model: [fallback_model]}] - with patch( - "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock - ): + with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock): with patch.object( processor, "common_processing_pre_call_logic", @@ -6312,9 +6201,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Sanity-check the premise: the limiter genuinely raises a # ProxyRateLimitError for the capped primary under the frozen clock. - with patch( - "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock - ): + with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock): with pytest.raises(ProxyRateLimitError): await limiter.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -6675,16 +6562,12 @@ class TestStreamingClientDisconnectBilling: prompt_tokens=1000, completion_tokens=10, total_tokens=1010, - prompt_tokens_details=PromptTokensDetailsWrapper( - cached_tokens=500 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500), ), ) ) - event = await self._bill_and_collect_success_event( - append_openai_style_cached_usage_chunk - ) + event = await self._bill_and_collect_success_event(append_openai_style_cached_usage_chunk) usage = event["response_obj"].usage assert getattr(usage, "cache_read_input_tokens", None) == 500 @@ -7454,9 +7337,7 @@ class TestInjectCostIntoUsageDict: logging_obj.model_call_details["custom_llm_provider"] = "anthropic" assert logging_obj.cost_breakdown is None - model_response = ModelResponse( - usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) - ) + model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)) cost = ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) assert cost is not None and cost > 0 @@ -7485,9 +7366,7 @@ class TestInjectCostIntoUsageDict: ) existing = logging_obj.cost_breakdown - model_response = ModelResponse( - usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) - ) + model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)) ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) assert logging_obj.cost_breakdown is existing @@ -7782,9 +7661,7 @@ def test_ttft_keepalive_interval_only_arms_for_a_streaming_request(request_data, @pytest.mark.asyncio @pytest.mark.parametrize("stream_requested, expect_ping", [(True, True), (False, False)]) -async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running( - stream_requested, expect_ping -): +async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(stream_requested, expect_ping): """The wiring, not the helper: every route funnels through this method, and the whole time-to-first-token is spent inside the call it wraps.""" @@ -7930,9 +7807,7 @@ async def test_a_late_failure_is_reported_to_the_failure_hook(): async def record(exc): audited.append(exc) - response = await open_sse_before_first_byte( - slow_failure(), ping_interval_seconds=0.05, on_late_failure=record - ) + response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=record) collected = await _drain(response) assert [type(exc).__name__ for exc in audited] == ["HTTPException"] @@ -7949,9 +7824,7 @@ async def test_a_failing_audit_hook_never_costs_the_client_its_error_frame(): async def broken_hook(exc): raise RuntimeError("the audit backend is down") - response = await open_sse_before_first_byte( - slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook - ) + response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8005,9 +7878,7 @@ async def test_base_process_llm_request_audits_a_failure_that_lands_after_its_ke [(0, False), (None, True)], ids=["operator-hard-disabled-this-deployment", "deployment-says-nothing"], ) -async def test_base_process_llm_request_honours_a_deployment_hard_disable( - deployment_keepalive, expect_ping -): +async def test_base_process_llm_request_honours_a_deployment_hard_disable(deployment_keepalive, expect_ping): """`keepalive_seconds: 0` is documented as a disable a request cannot lift. The funnel has to hand its router to the gate for that to hold before the upstream has answered, since no deployment has served the request yet.""" @@ -8053,9 +7924,7 @@ async def test_a_hook_returning_a_replacement_decides_what_the_client_sees(): async def sanitize(exc): return HTTPException(status_code=502, detail="upstream unavailable") - response = await open_sse_before_first_byte( - slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize - ) + response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8094,9 +7963,7 @@ async def test_a_hook_that_returns_nothing_leaves_the_real_error_intact(): async def audit_only(exc): return None - response = await open_sse_before_first_byte( - slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only - ) + response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8113,9 +7980,7 @@ async def test_a_broken_hook_does_not_replace_the_real_error_with_its_own_bug(): async def broken_hook(exc): raise RuntimeError("the audit backend is down") - response = await open_sse_before_first_byte( - slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook - ) + response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8328,9 +8193,7 @@ class TestStreamingResponseHeadersFollowFallback: proxy_logging_obj.post_call_success_hook = AsyncMock( side_effect=lambda data, user_api_key_dict, response: response ) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock( - return_value={"x-callback-header": "kept"} - ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-callback-header": "kept"}) async def fake_route_request(**kwargs): async def call(): @@ -8338,9 +8201,7 @@ class TestStreamingResponseHeadersFollowFallback: return call() - monkeypatch.setattr( - litellm.proxy.common_request_processing, "route_request", fake_route_request - ) + monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) result = await processor.base_process_llm_request( request=Request(scope={"type": "http", "headers": []}), From 1b594fc93515dd70168325839679e1aad2df71be Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:24:03 -0700 Subject: [PATCH 085/107] fix(guardrails): scan empty top-level system text blocks too The hoisted structured row keeps every text block of the top-level system prompt, empty ones included, while the scanned texts dropped the empty ones. Guardrails that count one text per slot then came back with more texts than the handler could place, so their rewrite was rejected. User text blocks were already scanned empty or not; the system prompt now matches. --- .../chat/guardrail_translation/handler.py | 2 +- .../test_anthropic_guardrail_handler.py | 42 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 1f4e1487316..9f3c85b555a 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -721,7 +721,7 @@ class AnthropicMessagesHandler(BaseTranslation): return tuple( ScannedText(text_str, SystemBlockTextTarget(block_idx)) for block_idx, block in enumerate(content) - if isinstance(block, dict) and isinstance(text_str := block.get("text"), str) and text_str + if isinstance(block, dict) and isinstance(text_str := block.get("text"), str) ) @staticmethod diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 1cacbb6be6b..287674afaa7 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2499,6 +2499,29 @@ class PerRowTextGuardrail(CustomGuardrail): return {**inputs, "texts": [str(row.get("content")).replace("123-45-6789", "") for row in rows]} +class PerSlotTextGuardrail(CustomGuardrail): + """Answers one redacted text per text slot of every chat row it was shown, the + way a guardrail that counts slots per message does, and hands back only texts.""" + + def __init__(self): + super().__init__(guardrail_name="per-slot-redactor") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts + + rows = inputs.get("structured_messages") or [] + return { + **inputs, + "texts": [text.replace("123-45-6789", "") for row in rows for text in message_slot_texts(row)], + } + + class TestPerMessageTextWriteBack: """Texts that no longer pair one-to-one with what the handler extracted must be rejected by name instead of sliding onto the wrong messages.""" @@ -2537,6 +2560,25 @@ class TestPerMessageTextWriteBack: assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched" assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched" + @pytest.mark.asyncio + async def test_one_text_per_slot_over_a_system_prompt_with_an_empty_block_is_applied(self): + data = { + "model": "claude-sonnet-4-5", + "system": [ + {"type": "text", "text": ""}, + {"type": "text", "text": "Reply with exactly the SSN you were given."}, + ], + "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}], + } + + await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerSlotTextGuardrail()) + + assert data["system"] == [ + {"type": "text", "text": ""}, + {"type": "text", "text": "Reply with exactly the SSN you were given."}, + ] + assert data["messages"] == [{"role": "user", "content": "My SSN is ."}] + @pytest.mark.asyncio async def test_one_text_per_row_without_a_system_prompt_is_applied(self): data = { From 80d804d6f98b69ff280b2020df76cc3bd6aa07fe Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 22:26:20 -0700 Subject: [PATCH 086/107] test(spend): preserve multi-day coverage and immutable assertions --- .../spend_tracking/spend_reconciliation.py | 12 +++-- .../spend_tracking/test_spend_tracking_e2e.py | 13 +++-- .../test_team_daily_activity_e2e.py | 48 +++++++++++++------ 3 files changed, 51 insertions(+), 22 deletions(-) diff --git a/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py index 8fcfea3f296..26809874aed 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py +++ b/tests/e2e/quota_management/spend_tracking/spend_reconciliation.py @@ -95,9 +95,10 @@ def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None: assert frozenset(row.request_id for row in rows) == expected_ids, "stored IDs must equal returned response IDs" assert len(rows) == len(traffic.responses), "expected exactly one scoped spend row per response" by_id: Final = {row.request_id: row for row in rows} - for response in traffic.responses: - row = by_id[response.id] - usage = response.usage + + def assert_response(response: ChatResponse) -> None: + row: Final = by_id[response.id] + usage: Final = response.usage assert usage is not None and usage.prompt_tokens is not None and usage.completion_tokens is not None assert row.team_id == traffic.team_id assert row.status == "success" @@ -105,5 +106,8 @@ def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None: assert row.prompt_tokens == usage.prompt_tokens assert row.completion_tokens == usage.completion_tokens assert row.total_tokens == usage.total_tokens - expected_cost = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE + expected_cost: Final = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE assert row.spend is not None and isclose(row.spend, expected_cost, rel_tol=1e-6, abs_tol=1e-9) + + for response in traffic.responses: + assert_response(response) diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py index 750285cbfb6..8a91e53e7d7 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_tracking_e2e.py @@ -18,6 +18,7 @@ fails the test; a pricing or token-count drift does not. import time from collections.abc import Callable from math import isclose +from typing import Final import pytest from e2e_http import Success @@ -284,14 +285,18 @@ def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> N def test_burst_of_concurrent_calls_loses_no_spend( client: SpendClient, resources: ResourceManager ) -> None: - from spend_reconciliation import assert_logs_match, create_traffic + from spend_reconciliation import TeamTraffic, assert_logs_match, create_traffic - traffic = create_traffic(client, resources) - for team in traffic: + traffic: Final = create_traffic(client, resources) + + def assert_team(team: TeamTraffic) -> None: assert_logs_match(client, team) - key_spend = client.poll_key_spend(team.key, minimum=team.spend * 0.999999) + key_spend: Final = client.poll_key_spend(team.key, minimum=team.spend * 0.999999) assert isclose(key_spend, team.spend, rel_tol=1e-6, abs_tol=1e-9) + for team in traffic: + assert_team(team) + @pytest.mark.covers("quota_management.spend_tracking.pagination.keeps_total") def test_spend_logs_v2_pagination_caps_pages_and_keeps_total( diff --git a/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py b/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py index dca5572f510..ef635e59743 100644 --- a/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_team_daily_activity_e2e.py @@ -15,9 +15,10 @@ from typing import Final import pytest from e2e_http import ProbeResult from lifecycle import ResourceManager +from proxy_client import Converged, await_converged from pydantic import BaseModel from spend_e2e_client import SpendClient -from spend_reconciliation import assert_logs_match, create_traffic +from spend_reconciliation import TeamTraffic, assert_logs_match, create_traffic pytestmark = pytest.mark.e2e @@ -92,7 +93,7 @@ class TestTeamDailyActivity: team_ids: Final = ",".join(team.team_id for team in traffic) def fetch( - page: int, start: str = started.isoformat(), end: str = ended.isoformat() + page: int, start: str = (started - timedelta(days=1)).isoformat(), end: str = ended.isoformat() ) -> TeamDailyActivityResponse: result: Final = _probe( client, @@ -112,22 +113,27 @@ class TestTeamDailyActivity: assert first.metadata.total_pages <= len(traffic) * 2, "unexpected extra scoped daily groups" return (first, *(fetch(page) for page in range(2, first.metadata.total_pages + 1))) - deadline: Final = time.monotonic() + client.proxy.poll_timeout - while True: - observed = pages() - if sum(page.metadata.total_api_requests for page in observed) >= sum(len(t.responses) for t in traffic): - break - if time.monotonic() >= deadline: - break - time.sleep(client.proxy.poll_interval) + outcome: Final = await_converged( + pages, + converged=lambda values: ( + sum(page.metadata.total_api_requests for page in values) >= sum(len(team.responses) for team in traffic) + ), + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + observed: Final = outcome.result if isinstance(outcome, Converged) else outcome.last_result + assert observed is not None, "daily aggregation must return a response before the deadline" assert len(observed) >= 2, "two teams must exercise a page boundary" - for index, page in enumerate(observed, 1): + + def assert_page(index: int, page: TeamDailyActivityResponse) -> None: assert page.metadata.page == index assert page.metadata.total_pages == len(observed) assert page.metadata.has_more == (index < len(observed)) assert len(page.results) == 1, "each fetched daily group must appear in results" - row = page.results[0] + row: Final = page.results[0] assert started <= datetime.fromisoformat(row.date).date() <= ended assert len(row.breakdown.entities) == 1 assert row.metrics.total_tokens == page.metadata.total_tokens @@ -138,6 +144,9 @@ class TestTeamDailyActivity: assert row.metrics.failed_requests == page.metadata.total_failed_requests assert isclose(row.metrics.spend, page.metadata.total_spend, rel_tol=1e-6, abs_tol=1e-9) + for index, page in enumerate(observed, 1): + assert_page(index, page) + entities: Final = tuple( (team_id, entity.metrics) for page in observed @@ -145,8 +154,9 @@ class TestTeamDailyActivity: for team_id, entity in row.breakdown.entities.items() ) assert frozenset(team_id for team_id, _ in entities) == frozenset(team.team_id for team in traffic) - for team in traffic: - metrics = tuple(metrics for team_id, metrics in entities if team_id == team.team_id) + + def assert_team(team: TeamTraffic) -> None: + metrics: Final = tuple(metrics for team_id, metrics in entities if team_id == team.team_id) assert sum(m.api_requests for m in metrics) == len(team.responses) assert sum(m.successful_requests for m in metrics) == len(team.responses) assert sum(m.failed_requests for m in metrics) == 0 @@ -154,6 +164,10 @@ class TestTeamDailyActivity: assert sum(m.completion_tokens for m in metrics) == team.completion_tokens assert sum(m.total_tokens for m in metrics) == team.prompt_tokens + team.completion_tokens assert isclose(sum(m.spend for m in metrics), team.spend, rel_tol=1e-6, abs_tol=1e-9) + + for team in traffic: + assert_team(team) + assert isclose( sum(page.metadata.total_spend for page in observed), sum(team.spend for team in traffic), @@ -164,6 +178,12 @@ class TestTeamDailyActivity: team.prompt_tokens + team.completion_tokens for team in traffic ) + for days in (7, 30): + assert ( + tuple(fetch(page, (started - timedelta(days=days)).isoformat()) for page in range(1, len(observed) + 1)) + == observed + ), f"{days}-day activity must preserve the same isolated groups and totals" + empty_date: Final = (started - timedelta(days=7)).isoformat() empty: Final = fetch(1, empty_date, empty_date) assert empty.results == [] From 6764ab2673942e7a32a3a12414c5e7ecec64a8fa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:29:32 -0700 Subject: [PATCH 087/107] test(router): assert num_retries_per_request as a per-group cap that resets per fallback hop #40930 (LIT-7505) changed num_retries_per_request from a request-wide cap to a per-model-group cap that resets on every fallback hop, and its own comment in litellm/__init__.py names that contract. The legacy test_async_fallbacks_max_retries_per_request still asserted the old request-wide reading (previous_models == 0), so the CircleCI router suite has been red on main since that merge for every run-ci PR. The test now reads the flat RetryAttemptRecord entries the fallback call carries and asserts the new contract directly: every record is from the first group, the retry at attempted_retries 0 is the real AuthenticationError, and each later attempt was refused with "Max retries per request hit!". --- tests/local_testing/test_router_fallbacks.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 82b832f89fd..5d7955ad34a 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -5,6 +5,7 @@ import asyncio import os import time import traceback +from typing import Final import pytest @@ -13,6 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger +from litellm.types.router import RetryAttemptRecord from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE @@ -21,15 +23,15 @@ class MyCustomHandler(CustomLogger): success: bool = False failure: bool = False previous_models: int = 0 + previous_model_records: tuple[RetryAttemptRecord, ...] = () def log_pre_api_call(self, model, messages, kwargs): print(f"Pre-API Call") print( f"previous_models: {kwargs['litellm_params']['metadata'].get('previous_models', None)}" ) - self.previous_models = len( - kwargs["litellm_params"]["metadata"].get("previous_models", []) - ) # {"previous_models": [{"model": litellm_model_name, "exception_type": AuthenticationError, "exception_string": }]} + self.previous_model_records = tuple(kwargs["litellm_params"]["metadata"].get("previous_models", ())) + self.previous_models = len(self.previous_model_records) print(f"self.previous_models: {self.previous_models}") def log_post_api_call(self, kwargs, response_obj, start_time, end_time): @@ -718,7 +720,14 @@ async def test_async_fallbacks_max_retries_per_request(): await asyncio.sleep( 0.05 ) # allow a delay as success_callbacks are on a separate thread - assert customHandler.previous_models == 0 # 0 retries, 0 fallback + records: Final = customHandler.previous_model_records + assert customHandler.previous_models == len(records) + assert records + assert {record["model_group"] for record in records} == {"azure/gpt-3.5-turbo"} + assert next(record["exception_type"] for record in records if record["attempted_retries"] == 0) == "AuthenticationError" + refused_retries: Final = tuple(record for record in records if record["attempted_retries"]) + assert refused_retries + assert all("Max retries per request hit!" in record["exception_string"] for record in refused_retries) router.reset() except litellm.Timeout as e: pass From 62b2b36ce90e6054a77b5c648c4013bce09c271f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:39:19 -0700 Subject: [PATCH 088/107] test(proxy): drop the reformat-only diff of the request processing tests The proxy edge test file no longer carries any test of this change, and the remaining diff was the scoped format gate reflowing the whole file to the 120 limit, so it goes back to the merge base bytes --- .../proxy/test_common_request_processing.py | 443 +++++++++++------- 1 file changed, 285 insertions(+), 158 deletions(-) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index a33b491fce2..812fd8ed47d 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -127,12 +127,16 @@ class TestProxyBaseLLMRequestProcessing: assert json.loads(result.body) == guardrailed_body @pytest.mark.asyncio - async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch): + async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers( + self, monkeypatch + ): """The guardrail JSON path must forward upstream response headers (e.g. x-amzn-requestid) alongside the x-litellm-* headers, matching the non-guardrail passthrough path, while dropping length headers that no longer match the rewritten body.""" - processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"custom_llm_provider": "bedrock"} + ) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -172,10 +176,14 @@ class TestProxyBaseLLMRequestProcessing: assert result.headers["content-length"] == str(len(result.body)) @pytest.mark.asyncio - async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch): + async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers( + self, monkeypatch + ): """The guardrail event-stream branch must also forward upstream response headers alongside the x-litellm-* headers.""" - processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"custom_llm_provider": "bedrock"} + ) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -217,11 +225,15 @@ class TestProxyBaseLLMRequestProcessing: assert result.headers["x-litellm-call-id"] == "test-call-id" @pytest.mark.asyncio - async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(self, monkeypatch): + async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook( + self, monkeypatch + ): """Guardrailed non-streaming passthrough responses must include headers injected by post_call_response_headers_hook, matching the headers a non-guardrailed passthrough response would carry.""" - processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"custom_llm_provider": "bedrock"} + ) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -240,7 +252,9 @@ class TestProxyBaseLLMRequestProcessing: return kwargs["response"] proxy_logging_obj.post_call_success_hook = fake_post_call_success_hook - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-litellm-custom": "from-hook"}) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-litellm-custom": "from-hook"} + ) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=upstream, @@ -364,7 +378,9 @@ class TestProxyBaseLLMRequestProcessing: json.dumps(persisted_body) @pytest.mark.asyncio - async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails(self, monkeypatch): + async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails( + self, monkeypatch + ): """arm_pre_call must run before pre_call_hook: an auto router's own compression policy has to be in `data["metadata"]` (naming the model-side guardrail so it runs even if it isn't default_on) by the time guardrails see the request.""" @@ -2175,10 +2191,16 @@ class TestCommonRequestProcessingHelpers: def _stringified_none_paths(node: object, path: str = "error") -> tuple[str, ...]: if isinstance(node, dict): - return tuple(found for key, value in node.items() for found in _stringified_none_paths(value, f"{path}.{key}")) + return tuple( + found + for key, value in node.items() + for found in _stringified_none_paths(value, f"{path}.{key}") + ) if isinstance(node, (list, tuple)): return tuple( - found for index, value in enumerate(node) for found in _stringified_none_paths(value, f"{path}[{index}]") + found + for index, value in enumerate(node) + for found in _stringified_none_paths(value, f"{path}[{index}]") ) return (path,) if node == "None" else () @@ -2925,7 +2947,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), ) assert headers["x-litellm-response-duration-ms"] == "500.0" @@ -2944,7 +2968,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), read_timing_from_logging_obj=False, ) @@ -2965,7 +2991,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={"_response_ms": 300.0}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), ) assert headers["x-litellm-response-duration-ms"] == "300.0" @@ -3006,7 +3034,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={"_response_ms": 300.0, "litellm_overhead_time_ms": 7.5}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), ) assert headers["x-litellm-response-duration-ms"] == "300.0" @@ -3458,7 +3488,9 @@ class TestStreamCloseOnDisconnect: finally: closed.set() - response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream") + response = _UpstreamClosingStreamingResponse( + body(), media_type="text/event-stream" + ) async def receive(): await asyncio.Event().wait() @@ -3489,7 +3521,9 @@ class TestStreamCloseOnDisconnect: finally: closed.set() - response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream") + response = _UpstreamClosingStreamingResponse( + body(), media_type="text/event-stream" + ) async def receive(): await disconnected.wait() @@ -3560,7 +3594,9 @@ class TestStreamCloseOnDisconnect: finally: inner_closed.set() - response = await create_response(generator=wrapped(), media_type="text/event-stream", headers={}) + response = await create_response( + generator=wrapped(), media_type="text/event-stream", headers={} + ) async def receive(): await asyncio.Event().wait() @@ -3792,7 +3828,9 @@ class TestStreamCloseOnDisconnect: with pytest.raises(_ClientDisconnectedBeforeFirstChunk): await asyncio.wait_for( - _buffer_first_chunk_honoring_disconnect(AcloseRaises(), request=self._request_that_disconnects()), + _buffer_first_chunk_honoring_disconnect( + AcloseRaises(), request=self._request_that_disconnects() + ), timeout=5, ) @@ -3808,7 +3846,9 @@ class TestStreamCloseOnDisconnect: with pytest.raises(_ClientDisconnectedBeforeFirstChunk): await asyncio.wait_for( - _buffer_first_chunk_honoring_disconnect(blocking_gen(), request=self._request_that_disconnects()), + _buffer_first_chunk_honoring_disconnect( + blocking_gen(), request=self._request_that_disconnects() + ), timeout=5, ) assert closed.is_set() @@ -3824,7 +3864,9 @@ class TestHandleLLMApiExceptionRetryAfter: user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=callback_headers or {}) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value=callback_headers or {} + ) try: await processor._handle_llm_api_exception( @@ -3876,7 +3918,9 @@ class TestHandleLLMApiExceptionRetryAfter: enable_pre_call_checks=False, cooldown_list=[], ) - proxy_exc = await self._invoke(exc, callback_headers={"retry-after": "", "x-custom": "1"}) + proxy_exc = await self._invoke( + exc, callback_headers={"retry-after": "", "x-custom": "1"} + ) assert proxy_exc.headers["retry-after"] == "43" assert proxy_exc.headers["x-custom"] == "1" @@ -4105,7 +4149,9 @@ class TestDisconnectGatherCleanup: return Request(scope={"type": "http", "headers": []}, receive=receive) @pytest.mark.asyncio - async def test_base_process_llm_request_raises_499_on_client_disconnect(self, monkeypatch): + async def test_base_process_llm_request_raises_499_on_client_disconnect( + self, monkeypatch + ): """With cancel_on_disconnect enabled, base_process_llm_request returns 499.""" import asyncio @@ -4134,7 +4180,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) with pytest.raises(HTTPException) as exc_info: await processing_obj.base_process_llm_request( @@ -4152,7 +4200,9 @@ class TestDisconnectGatherCleanup: assert "disconnected" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(self, monkeypatch): + async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect( + self, monkeypatch + ): import asyncio import litellm.proxy.common_request_processing as cpr @@ -4177,7 +4227,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) monkeypatch.setattr( cpr, "route_request", @@ -4238,7 +4290,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) with pytest.raises(HTTPException): await processing_obj.base_process_llm_request( @@ -4289,7 +4343,9 @@ class TestDisconnectGatherCleanup: assert task.done() @pytest.mark.asyncio - async def test_base_process_llm_request_preserves_llm_error_after_gather(self, monkeypatch): + async def test_base_process_llm_request_preserves_llm_error_after_gather( + self, monkeypatch + ): import litellm.proxy.common_request_processing as cpr from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -4318,7 +4374,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) mock_request = MagicMock(spec=Request) mock_request.is_disconnected = AsyncMock(return_value=False) @@ -4355,13 +4413,19 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": {}}, } - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert request_data["metadata"]["error_information"]["error_code"] == "499" assert ( - mock_logging_obj.model_call_details["litellm_params"]["metadata"]["error_information"]["error_code"] + request_data["metadata"]["error_information"]["error_code"] == "499" + ) + assert ( + mock_logging_obj.model_call_details["litellm_params"]["metadata"][ + "error_information" + ]["error_code"] == "499" ) @@ -4375,7 +4439,9 @@ class TestStreamingClientDisconnectLogging: mock_request.is_disconnected = AsyncMock(return_value=False) request_data = {"metadata": {}} - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is False assert "client_disconnected" not in request_data["metadata"] @@ -4400,12 +4466,22 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": {}}, } - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert mock_logging_obj.model_call_details["litellm_params"]["metadata"]["client_disconnected"] is True - assert mock_logging_obj.model_call_details["metadata"]["client_disconnected"] is True + assert ( + mock_logging_obj.model_call_details["litellm_params"]["metadata"][ + "client_disconnected" + ] + is True + ) + assert ( + mock_logging_obj.model_call_details["metadata"]["client_disconnected"] + is True + ) @pytest.mark.asyncio async def test_record_streaming_client_disconnect_handles_none_request_data_metadata(self): @@ -4421,11 +4497,15 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": None}, } - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert request_data["litellm_params"]["metadata"]["client_disconnected"] is True + assert ( + request_data["litellm_params"]["metadata"]["client_disconnected"] is True + ) @pytest.mark.asyncio async def test_apply_client_disconnect_metadata_none_returns_early(self): @@ -4436,7 +4516,9 @@ class TestStreamingClientDisconnectLogging: _apply_client_disconnect_metadata(None) @pytest.mark.asyncio - async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(self, monkeypatch): + async def test_finalize_streaming_generator_cleanup_fires_deferred_logging( + self, monkeypatch + ): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4468,7 +4550,9 @@ class TestStreamingClientDisconnectLogging: assert request_data["metadata"]["error_information"]["error_code"] == "499" @pytest.mark.asyncio - async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(self, monkeypatch): + async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion( + self, monkeypatch + ): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4498,7 +4582,9 @@ class TestStreamingClientDisconnectLogging: assert "client_disconnected" not in request_data["metadata"] @pytest.mark.asyncio - async def test_async_streaming_data_generator_records_499_on_early_aclose(self, monkeypatch): + async def test_async_streaming_data_generator_records_499_on_early_aclose( + self, monkeypatch + ): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4513,7 +4599,9 @@ class TestStreamingClientDisconnectLogging: yield {"choices": [{"delta": {"content": " there"}}]} mock_proxy_logging = MagicMock(spec=ProxyLogging) - mock_proxy_logging.async_post_call_streaming_iterator_hook = mock_streaming_iterator + mock_proxy_logging.async_post_call_streaming_iterator_hook = ( + mock_streaming_iterator + ) ProxyLogging._callback_capabilities_cache.clear() mock_request = MagicMock(spec=Request) @@ -4524,7 +4612,9 @@ class TestStreamingClientDisconnectLogging: "model": "gemini-2.0-flash", "metadata": {}, "litellm_params": {"metadata": {}}, - "litellm_logging_obj": MagicMock(model_call_details={"metadata": {}, "litellm_params": {}}), + "litellm_logging_obj": MagicMock( + model_call_details={"metadata": {}, "litellm_params": {}} + ), } gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( @@ -4543,8 +4633,6 @@ class TestStreamingClientDisconnectLogging: assert request_data["metadata"]["error_information"]["error_code"] == "499" ProxyLogging._callback_capabilities_cache.clear() - - class TestCancelOnDisconnect: """ Coverage for the opt-in `general_settings.cancel_on_disconnect` flag: @@ -4571,17 +4659,23 @@ class TestCancelOnDisconnect: llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) + await _cancel_llm_call_on_client_disconnect( + request, llm_call, disconnect_event + ) assert llm_call.cancelled() assert disconnect_event.is_set() async def test_monitor_is_noop_while_client_stays_connected(self): - request = self._request([{"type": "http.request", "body": b"", "more_body": False}]) + request = self._request( + [{"type": "http.request", "body": b"", "more_body": False}] + ) llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)) + monitor = asyncio.create_task( + _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) + ) await asyncio.sleep(0.01) assert not monitor.done() @@ -4600,7 +4694,9 @@ class TestCancelOnDisconnect: llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) + await _cancel_llm_call_on_client_disconnect( + request, llm_call, disconnect_event + ) assert not llm_call.cancelled() assert not disconnect_event.is_set() @@ -4615,7 +4711,9 @@ class TestCancelOnDisconnect: with pytest.raises(asyncio.CancelledError): await _await_llm_call_cancelling_on_disconnect(request, llm_call) - async def _drive_base_process_llm_request(self, monkeypatch, general_settings: dict, llm_call, request: Request): + async def _drive_base_process_llm_request( + self, monkeypatch, general_settings: dict, llm_call, request: Request + ): from litellm.proxy._types import UserAPIKeyAuth logging_obj = MagicMock() @@ -4624,7 +4722,9 @@ class TestCancelOnDisconnect: logging_obj._on_deferred_stream_complete = None logging_obj.cost_breakdown = None - processor = ProxyBaseLLMRequestProcessing(data={"model": "fake-model", "litellm_logging_obj": logging_obj}) + processor = ProxyBaseLLMRequestProcessing( + data={"model": "fake-model", "litellm_logging_obj": logging_obj} + ) proxy_logging_obj = MagicMock(spec=ProxyLogging) proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) @@ -4632,7 +4732,9 @@ class TestCancelOnDisconnect: proxy_logging_obj.post_call_success_hook = AsyncMock( side_effect=lambda data, user_api_key_dict, response: response ) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value=None + ) async def fake_route_request(**kwargs): return llm_call() @@ -4711,7 +4813,9 @@ class TestCancelOnDisconnect: with pytest.raises(ProxyException) as exc_info: await processor._handle_llm_api_exception( - e=HTTPException(status_code=499, detail="Client disconnected the request"), + e=HTTPException( + status_code=499, detail="Client disconnected the request" + ), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), proxy_logging_obj=proxy_logging_obj, ) @@ -4777,9 +4881,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", capture_hook) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4829,9 +4931,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", non_dict_hook) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4869,9 +4969,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: hook_spy = AsyncMock() monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4912,9 +5010,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: hook_spy = AsyncMock() monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -5026,9 +5122,7 @@ class TestEventStreamAllmPassthroughRoute: "content-length": "99", } - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=mock_response, @@ -5059,7 +5153,9 @@ class TestAllmPassthroughStreamingProviderGate: de-anonymized. """ - def _build_processing_obj(self, custom_llm_provider: str, endpoint: str = "") -> ProxyBaseLLMRequestProcessing: + def _build_processing_obj( + self, custom_llm_provider: str, endpoint: str = "" + ) -> ProxyBaseLLMRequestProcessing: logging_obj = MagicMock() logging_obj.litellm_call_id = "call-123" logging_obj.cost_breakdown = None @@ -5146,17 +5242,14 @@ class TestAllmPassthroughStreamingProviderGate: processing_obj = self._build_processing_obj("anthropic") chunks = [b"chunk-1", b"chunk-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5165,27 +5258,27 @@ class TestAllmPassthroughStreamingProviderGate: assert streamed == chunks @pytest.mark.asyncio - async def test_bedrock_converse_stream_is_buffered_through_handler(self, monkeypatch): - processing_obj = self._build_processing_obj("bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream") + async def test_bedrock_converse_stream_is_buffered_through_handler( + self, monkeypatch + ): + processing_obj = self._build_processing_obj( + "bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream" + ) chunks = [b"raw-1", b"raw-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), - patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=b"modified-body"), - ) as mock_handler, - ): + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler: result = await self._run(processing_obj, monkeypatch, chunks) assert isinstance(result, Response) @@ -5201,23 +5294,19 @@ class TestAllmPassthroughStreamingProviderGate: ) chunks = [b"raw-1", b"raw-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), - patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=b"modified-body"), - ) as mock_handler, - ): + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler: result = await self._run(processing_obj, monkeypatch, chunks) assert isinstance(result, StreamingResponse) @@ -5239,17 +5328,14 @@ class TestAllmPassthroughStreamingProviderGate: ) chunks = [b"raw-1", b"raw-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=False, - ), + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5268,17 +5354,14 @@ class TestAllmPassthroughStreamingProviderGate: processing_obj = self._build_processing_obj("anthropic") chunks = [b"chunk-1", b"chunk-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=False, - ), + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5726,7 +5809,9 @@ class TestCostHeadersForCallsPricedAtZero: fastapi_response = Response() processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj}) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False + ): await processing_obj.base_process_llm_request( request=MagicMock(spec=Request, headers={}), fastapi_response=fastapi_response, @@ -5797,7 +5882,9 @@ class TestCostHeadersForCallsPricedAtZero: assert breakdown.tool_usage_cost == 0.0 def test_cost_breakdown_stays_empty_for_an_inference_call(self): - breakdown = _get_cost_breakdown_from_logging_obj(litellm_logging_obj=self._logging_obj(call_type="acompletion")) + breakdown = _get_cost_breakdown_from_logging_obj( + litellm_logging_obj=self._logging_obj(call_type="acompletion") + ) assert breakdown == CostBreakdownHeaderValues() @@ -5824,6 +5911,7 @@ class TestCostHeadersForCallsPricedAtZero: class TestPreCallWithFallbacksOnLocalRateLimit: + @pytest.mark.asyncio async def test_fallback_triggered_on_local_rate_limit(self): from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -5975,7 +6063,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] user_api_key_dict = MagicMock() - user_api_key_dict.router_settings = {"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]} + user_api_key_dict.router_settings = { + "fallbacks": [{"gpt-4": ["claude-3-haiku"]}] + } with patch.object( processor, @@ -6006,7 +6096,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "disable_fallbacks": True}) + processor = ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4", "disable_fallbacks": True} + ) async def mock_pre_call_logic(**kwargs): raise ProxyRateLimitError( @@ -6132,7 +6224,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Real per-key per-model TPM limiter + a key carrying the customer's # `model_tpm_limit` metadata (only the primary is capped). - limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + limiter = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) user_api_key_dict = UserAPIKeyAuth( api_key="sk-lit3890", metadata={"model_tpm_limit": {primary_model: 100}}, @@ -6140,7 +6234,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Pre-seed the primary's per-model token counter at the cap so the very # next request trips it. The counter key uses the *hashed* api_key. - counter_key = f"{user_api_key_dict.api_key}::{primary_model}::{precise_minute}::request_count" + counter_key = ( + f"{user_api_key_dict.api_key}::{primary_model}" + f"::{precise_minute}::request_count" + ) await limiter.internal_usage_cache.async_set_cache( key=counter_key, value={"current_requests": 0, "current_tpm": 100, "current_rpm": 0}, @@ -6171,7 +6268,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router = MagicMock() mock_router.fallbacks = [{primary_model: [fallback_model]}] - with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock): + with patch( + "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock + ): with patch.object( processor, "common_processing_pre_call_logic", @@ -6201,7 +6300,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Sanity-check the premise: the limiter genuinely raises a # ProxyRateLimitError for the capped primary under the frozen clock. - with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock): + with patch( + "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock + ): with pytest.raises(ProxyRateLimitError): await limiter.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -6562,12 +6663,16 @@ class TestStreamingClientDisconnectBilling: prompt_tokens=1000, completion_tokens=10, total_tokens=1010, - prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500), + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=500 + ), ), ) ) - event = await self._bill_and_collect_success_event(append_openai_style_cached_usage_chunk) + event = await self._bill_and_collect_success_event( + append_openai_style_cached_usage_chunk + ) usage = event["response_obj"].usage assert getattr(usage, "cache_read_input_tokens", None) == 500 @@ -7337,7 +7442,9 @@ class TestInjectCostIntoUsageDict: logging_obj.model_call_details["custom_llm_provider"] = "anthropic" assert logging_obj.cost_breakdown is None - model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)) + model_response = ModelResponse( + usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) + ) cost = ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) assert cost is not None and cost > 0 @@ -7366,7 +7473,9 @@ class TestInjectCostIntoUsageDict: ) existing = logging_obj.cost_breakdown - model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)) + model_response = ModelResponse( + usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) + ) ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) assert logging_obj.cost_breakdown is existing @@ -7661,7 +7770,9 @@ def test_ttft_keepalive_interval_only_arms_for_a_streaming_request(request_data, @pytest.mark.asyncio @pytest.mark.parametrize("stream_requested, expect_ping", [(True, True), (False, False)]) -async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(stream_requested, expect_ping): +async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running( + stream_requested, expect_ping +): """The wiring, not the helper: every route funnels through this method, and the whole time-to-first-token is spent inside the call it wraps.""" @@ -7807,7 +7918,9 @@ async def test_a_late_failure_is_reported_to_the_failure_hook(): async def record(exc): audited.append(exc) - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=record) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=record + ) collected = await _drain(response) assert [type(exc).__name__ for exc in audited] == ["HTTPException"] @@ -7824,7 +7937,9 @@ async def test_a_failing_audit_hook_never_costs_the_client_its_error_frame(): async def broken_hook(exc): raise RuntimeError("the audit backend is down") - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -7878,7 +7993,9 @@ async def test_base_process_llm_request_audits_a_failure_that_lands_after_its_ke [(0, False), (None, True)], ids=["operator-hard-disabled-this-deployment", "deployment-says-nothing"], ) -async def test_base_process_llm_request_honours_a_deployment_hard_disable(deployment_keepalive, expect_ping): +async def test_base_process_llm_request_honours_a_deployment_hard_disable( + deployment_keepalive, expect_ping +): """`keepalive_seconds: 0` is documented as a disable a request cannot lift. The funnel has to hand its router to the gate for that to hold before the upstream has answered, since no deployment has served the request yet.""" @@ -7924,7 +8041,9 @@ async def test_a_hook_returning_a_replacement_decides_what_the_client_sees(): async def sanitize(exc): return HTTPException(status_code=502, detail="upstream unavailable") - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -7963,7 +8082,9 @@ async def test_a_hook_that_returns_nothing_leaves_the_real_error_intact(): async def audit_only(exc): return None - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -7980,7 +8101,9 @@ async def test_a_broken_hook_does_not_replace_the_real_error_with_its_own_bug(): async def broken_hook(exc): raise RuntimeError("the audit backend is down") - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8193,7 +8316,9 @@ class TestStreamingResponseHeadersFollowFallback: proxy_logging_obj.post_call_success_hook = AsyncMock( side_effect=lambda data, user_api_key_dict, response: response ) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-callback-header": "kept"}) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-header": "kept"} + ) async def fake_route_request(**kwargs): async def call(): @@ -8201,7 +8326,9 @@ class TestStreamingResponseHeadersFollowFallback: return call() - monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) + monkeypatch.setattr( + litellm.proxy.common_request_processing, "route_request", fake_route_request + ) result = await processor.base_process_llm_request( request=Request(scope={"type": "http", "headers": []}), From 7a7770db0d54e4734175fd7e798eb389ff92a1bd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 14 Sep 2026 22:46:55 -0700 Subject: [PATCH 089/107] test(e2e): verify streamed answers and tool continuation --- tests/e2e/e2e_http.py | 2 + .../test_chat_stream_contract_e2e.py | 83 +++++++--- .../e2e/llm_translation/test_messages_e2e.py | 144 +++++++++++++++++- tests/e2e/models.py | 11 ++ .../test_streaming_iterator_tool_args.py | 8 +- 5 files changed, 218 insertions(+), 30 deletions(-) diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 67370c98274..e0a20495964 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -170,6 +170,7 @@ class StreamingResponse(BaseModel): # the consumed body is elided, so this is the only place they surface. stream_error: str | None = None stream_done: bool = False + stream_done_positions: tuple[int, ...] = () @property def ok(self) -> bool: @@ -647,6 +648,7 @@ def streaming_outcome( stream_events=[payload for payload, _ in events], stream_event_arrivals=[arrived for _, arrived in events], stream_done=any(payload == _SSE_DONE for payload, _ in payloads), + stream_done_positions=tuple(index for index, (payload, _) in enumerate(payloads) if payload == _SSE_DONE), stream_error=next( (line.decode(errors="replace")[:300] for line, _ in stamped if _is_stream_error_line(line)), None, diff --git a/tests/e2e/llm_translation/test_chat_stream_contract_e2e.py b/tests/e2e/llm_translation/test_chat_stream_contract_e2e.py index 4db2fe004c5..fdb76df703d 100644 --- a/tests/e2e/llm_translation/test_chat_stream_contract_e2e.py +++ b/tests/e2e/llm_translation/test_chat_stream_contract_e2e.py @@ -1,51 +1,90 @@ -"""Vendor §12.3: chat completions streaming SSE contract (LIT-4778). - -Asserts a streamed /chat/completions response is SSE, carries content chunks, -and terminates with the OpenAI [DONE] sentinel. -""" - from __future__ import annotations +from typing import Final + import pytest -from e2e_config import unique_marker +from e2e_config import provider_edge_base, unique_marker from e2e_http import require_successful_call from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, LiteLLMParamsBody +from models import ChatBody, ChatMessage, ChatStreamOptions, LiteLLMParamsBody, Usage from proxy_client import ProxyClient +from pydantic import BaseModel -pytestmark = pytest.mark.e2e +pytestmark = [pytest.mark.e2e, pytest.mark.replayable] + + +class _Delta(BaseModel): + content: str | None = None + + +class _Choice(BaseModel): + index: int + delta: _Delta + finish_reason: str | None = None + + +class _Chunk(BaseModel): + choices: tuple[_Choice, ...] + usage: Usage | None = None class TestChatStreamContract: @pytest.mark.covers("llm.chat_completions.openai.basic.stream.works") def test_chat_stream_is_sse_and_ends_with_done(self, proxy: ProxyClient, resources: ResourceManager) -> None: - model = f"e2e-chat-stream-{unique_marker()}" - model_id = proxy.create_model( + model: Final = f"e2e-chat-stream-{unique_marker()}" + base: Final = provider_edge_base("openai") + model_id: Final = proxy.create_model( model, - LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"), + LiteLLMParamsBody( + model="openai/gpt-5.6", + api_key="os.environ/OPENAI_API_KEY", + api_base=f"{base}/v1" if base else None, + ), ) resources.defer(lambda: proxy.delete_model(model_id)) - key = resources.key() - - result = proxy.chat_stream( + key: Final = resources.key() + expected: Final = "The amber kite crosses the quiet lake." + result: Final = proxy.chat_stream( key, ChatBody( model=model, messages=[ ChatMessage( - role="user", - content=f"Reply with the single word ok. {unique_marker()}", + role="user", content=f"Repeat exactly this sentence, with no additional text: {expected}" ) ], stream=True, - max_completion_tokens=32, - temperature=0.0, + stream_options=ChatStreamOptions(include_usage=True), + max_completion_tokens=256, + reasoning_effort="none", ), ) require_successful_call(result) assert result.is_streaming, f"expected SSE content-type, got {result.content_type!r}" assert result.stream_events, "stream returned no data events" - assert result.stream_done, ( - f"stream must terminate with [DONE]; " - f"chunks={result.chunks} done={result.stream_done} events={len(result.stream_events)}" + assert not result.stream_error, f"stream errored: {result.stream_error}" + assert result.stream_done, "stream must terminate with [DONE]" + assert result.stream_done_positions == (len(result.stream_events),), "[DONE] must occur once after all events" + chunks: Final = tuple(_Chunk.model_validate_json(event) for event in result.stream_events) + text_positions: Final = tuple( + i for i, chunk in enumerate(chunks) if any(c.delta.content for c in chunk.choices) ) + terminal_positions: Final = tuple( + i for i, chunk in enumerate(chunks) if any(c.finish_reason is not None for c in chunk.choices) + ) + assert text_positions, "stream completed without meaningful text" + assert len(terminal_positions) == 1, "expected exactly one terminal choice" + assert text_positions[0] < terminal_positions[0], "meaningful text must arrive before termination" + assert text_positions[-1] <= terminal_positions[0], "text arrived after termination" + assert all(c.index == 0 for chunk in chunks for c in chunk.choices) + assert tuple(c.finish_reason for c in chunks[terminal_positions[0]].choices) == ("stop",) + text: Final = "".join(c.delta.content or "" for chunk in chunks for c in chunk.choices) + assert text.strip() == expected, f"streamed answer was altered or incomplete: {text!r}" + usage_positions: Final = tuple(i for i, chunk in enumerate(chunks) if chunk.usage is not None) + assert usage_positions == (len(chunks) - 1,), "expected one final usage chunk" + assert terminal_positions[0] < usage_positions[0], "usage must follow the terminal choice" + usage: Final = chunks[-1].usage + assert usage is not None + assert usage.prompt_tokens is not None and usage.prompt_tokens > 0 + assert usage.completion_tokens is not None and usage.completion_tokens > 0 + assert usage.total_tokens == usage.prompt_tokens + usage.completion_tokens diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index c731b52acd8..ca58c30d40c 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -21,7 +21,12 @@ from e2e_http import assert_client_error, require_successful_call, unwrap from endpoints_client import EndpointsClient, MessagesResult from lifecycle import ResourceManager from models import ( + AnthropicAssistantTurn, + AnthropicContentBlock, AnthropicCustomTool, + AnthropicToolChoice, + AnthropicToolResultBlock, + AnthropicToolResultTurn, AnthropicMessagesBody, ChatMessage, JsonSchemaProperty, @@ -29,7 +34,7 @@ from models import ( SpendLogRow, ToolInputSchema, ) -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict pytestmark = [pytest.mark.e2e, pytest.mark.replayable] @@ -284,8 +289,139 @@ class TestAnthropicMessages: result = endpoints_client.proxy.transport.send( "/v1/messages", headers=endpoints_client.proxy.transport.bearer(key), - json=_OptionalMessagesBody( - messages=[ChatMessage(role="user", content="hi")], max_tokens=50 - ), + json=_OptionalMessagesBody(messages=[ChatMessage(role="user", content="hi")], max_tokens=50), ) assert_client_error(result, "messages missing model") + + +class _BridgeDelta(BaseModel): + type: str | None = None + partial_json: str | None = None + stop_reason: str | None = None + + +class _BridgeEvent(BaseModel): + type: str + index: int | None = None + content_block: AnthropicContentBlock | None = None + delta: _BridgeDelta | None = None + + +class _ParcelInput(BaseModel): + model_config = ConfigDict(extra="forbid", strict=True) + parcel: str + shelf: int + + +def _tool_from_stream(events: tuple[_BridgeEvent, ...]) -> AnthropicContentBlock: + starts: Final = tuple( + event + for event in events + if event.type == "content_block_start" + and event.content_block is not None + and event.content_block.type == "tool_use" + ) + assert len(starts) == 1, "expected exactly one tool call" + start: Final = starts[0] + block: Final = start.content_block + assert block is not None and block.id and start.index is not None + fragments: Final = tuple( + event + for event in events + if event.type == "content_block_delta" and event.delta is not None and event.delta.type == "input_json_delta" + ) + assert fragments, "tool stream contained no argument fragments" + assert all(event.index == start.index for event in fragments), "tool fragments changed index" + positions: Final = tuple(i for i, event in enumerate(events) if event in fragments) + stops: Final = tuple( + i for i, event in enumerate(events) if event.type == "content_block_stop" and event.index == start.index + ) + assert len(stops) == 1 and events.index(start) < positions[0] <= positions[-1] < stops[0] + assert tuple( + event.delta.stop_reason for event in events if event.type == "message_delta" and event.delta is not None + ) == ("tool_use",) + terminal_positions: Final = tuple(i for i, event in enumerate(events) if event.type == "message_delta") + assert len(terminal_positions) == 1 and stops[0] < terminal_positions[0] < len(events) - 1 + assert tuple(i for i, event in enumerate(events) if event.type == "message_stop") == (len(events) - 1,), ( + "tool stream did not terminate exactly once" + ) + arguments: Final = _ParcelInput.model_validate_json( + "".join(event.delta.partial_json or "" for event in fragments if event.delta is not None) + ) + return AnthropicContentBlock(type="tool_use", id=block.id, name=block.name, input=arguments.model_dump()) + + +def _parcel_result(tool: AnthropicContentBlock, result: AnthropicToolResultBlock) -> AnthropicToolResultTurn: + assert tool.id and result.tool_use_id == tool.id, "tool result ID does not match the emitted call" + return AnthropicToolResultTurn(content=[result]) + + +def _request_tool( + client: EndpointsClient, key: str, request: AnthropicMessagesBody, stream: bool +) -> AnthropicContentBlock: + if stream: + response: Final = client.proxy.messages_stream(key, request) + require_successful_call(response) + assert response.is_streaming and not response.stream_error + return _tool_from_stream(tuple(_BridgeEvent.model_validate_json(event) for event in response.stream_events)) + response_body: Final = unwrap(client.proxy.messages(key, request)) + blocks: Final = tuple(block for block in response_body.content or () if block.type == "tool_use") + assert len(blocks) == 1 + return blocks[0] + + +class TestOpenAIMessagesToolContinuation: + @pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"]) + def test_required_tool_arguments_and_correlated_result( + self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool + ) -> None: + model: Final = f"e2e-bridge-tool-{unique_marker()}" + base: Final = provider_edge_base("openai") + model_id: Final = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model="openai/gpt-5.6", api_key="os.environ/OPENAI_API_KEY", api_base=f"{base}/v1" if base else None + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key: Final = resources.key(models=[model]) + tool: Final = AnthropicCustomTool( + name="locate_parcel", + description="Look up the receipt for a parcel on a shelf. Return the receipt verbatim.", + input_schema=ToolInputSchema( + properties={"parcel": JsonSchemaProperty(type="string"), "shelf": JsonSchemaProperty(type="integer")}, + required=["parcel", "shelf"], + ), + ) + question: Final = ChatMessage( + role="user", + content="Call locate_parcel with parcel exactly amber-kite and shelf exactly 7. After the tool result, reply with only the receipt returned by the tool.", + ) + request: Final = AnthropicMessagesBody( + model=model, + max_tokens=2048, + messages=[question], + tools=[tool], + tool_choice=AnthropicToolChoice(type="tool", name=tool.name), + stream=stream, + ) + emitted: Final = _request_tool(endpoints_client, key, request, stream) + assert emitted.id and emitted.name == "locate_parcel" + assert emitted.input == {"parcel": "amber-kite", "shelf": 7}, "required tool arguments were lost or changed" + receipt: Final = f"receipt-{unique_marker()}" + result_turn: Final = _parcel_result(emitted, AnthropicToolResultBlock(tool_use_id=emitted.id, content=receipt)) + continuation: Final = unwrap( + endpoints_client.proxy.messages( + key, + AnthropicMessagesBody( + model=model, + max_tokens=2048, + tools=[tool], + tool_choice=AnthropicToolChoice(type="none"), + messages=[question, AnthropicAssistantTurn(content=[emitted]), result_turn], + ), + ) + ) + answer: Final = "".join(block.text or "" for block in continuation.content or ()) + assert answer.strip() == receipt, "continuation did not consume the correlated tool result" + assert all(block.type != "tool_use" for block in continuation.content or ()) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 6fca1268ebc..ba6d0fe3334 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -283,10 +283,15 @@ class ChatToolResultTurn(BaseModel): type ChatTurn = ChatMessage | ChatAssistantTurn | ChatToolResultTurn +class ChatStreamOptions(BaseModel): + include_usage: bool + + class ChatBody(BaseModel): model: str messages: Sequence[ChatTurn] stream: bool = False + stream_options: ChatStreamOptions | None = None max_tokens: int | None = None max_completion_tokens: int | None = None temperature: float | None = None @@ -488,12 +493,18 @@ class AnthropicToolResultTurn(BaseModel): type AnthropicMessage = ChatMessage | AnthropicAssistantTurn | AnthropicToolResultTurn +class AnthropicToolChoice(BaseModel): + type: Literal["auto", "any", "tool", "none"] + name: str | None = None + + class AnthropicMessagesBody(BaseModel): model: str messages: list[AnthropicMessage] max_tokens: int stream: bool | None = None tools: list[AnthropicTool] | None = None + tool_choice: AnthropicToolChoice | None = None guardrails: list[str] | None = None cache: dict[str, bool] | None = {"no-cache": True} diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py index 29e9279731d..a20aaf2e324 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_tool_args.py @@ -8,6 +8,8 @@ Without the fix, the AnthropicStreamWrapper silently dropped these arguments, causing tool_use blocks to arrive with empty input {}. """ +import json + from typing import List from unittest.mock import MagicMock @@ -139,9 +141,7 @@ async def test_async_stream_emits_input_json_delta_for_bundled_tool_args(): # Verify the delta carries the tool arguments delta_event = events[input_json_delta_idx] - assert delta_event["delta"][ - "partial_json" - ], "input_json_delta should have non-empty partial_json" + assert json.loads(delta_event["delta"]["partial_json"]) == {"location": "Boston"} @pytest.mark.asyncio @@ -300,7 +300,7 @@ def test_sync_stream_emits_input_json_delta_for_bundled_tool_args(): assert ( input_json_delta_idx == tool_start_idx + 1 ), "input_json_delta should immediately follow the tool_use content_block_start" - assert events[input_json_delta_idx]["delta"]["partial_json"] + assert json.loads(events[input_json_delta_idx]["delta"]["partial_json"]) == {"location": "Boston"} def test_sync_stream_no_extra_delta_when_tool_args_empty(): From e3152c011dfe23d26f55c31e7a7e5ce402859f32 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:58:25 -0700 Subject: [PATCH 090/107] fix(responses): classify streamed tool calls on the chat name and strip guardrail edits around the grammar block The streaming bridge restored the namespace before deciding whether a tool call was a custom tool, so a namespaced function sharing a short name with a nested custom tool streamed back as a custom_tool_call. Classify on the raw chat tool name first, the way the non-streaming path already does. The guardrail merge only stripped the namespace prefix and grammar suffix from the ends of the edited description, so a guardrail appending text after the grammar block left the block in the member description and the chat conversion appended it a second time. Strip the first occurrence of each instead. --- .../guardrail_translation/tool_merge.py | 2 +- .../streaming_iterator.py | 26 +++---- ...t_openai_responses_guardrail_tool_merge.py | 15 ++++ .../test_litellm_completion_responses.py | 77 +++++++++++++++++++ 4 files changed, 105 insertions(+), 15 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py index 0326e9b2bfd..ff67c6220e1 100644 --- a/litellm/llms/openai/responses/guardrail_translation/tool_merge.py +++ b/litellm/llms/openai/responses/guardrail_translation/tool_merge.py @@ -72,7 +72,7 @@ def _function_fields(tool: Tool) -> Tool: def _member_description(key: str, value: object, prefix: str, suffix: str) -> object: if key != "description" or not isinstance(value, str): return value - return value.removeprefix(prefix).removesuffix(suffix) + return value.replace(prefix, "", 1).replace(suffix, "", 1) def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool: diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 126b976e2c5..c28b5558c75 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -8,6 +8,7 @@ from litellm.main import stream_chunk_builder from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, extract_custom_tool_names, + is_custom_tool_call, serialize_tool_call_arguments, ) from litellm.responses.litellm_completion_transformation.transformation import ( @@ -166,6 +167,14 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return tool_name, namespace return fn_name, None + def _tool_call_item_kwargs(self, call_id: str, fn_name: str, arguments: str, status: str) -> dict[str, str]: + item_kwargs: Final = build_tool_call_item_kwargs(call_id, fn_name, arguments, status, self._custom_tool_names) + if is_custom_tool_call(fn_name, self._custom_tool_names): + return item_kwargs + tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) + namespace_kwargs: Final = {"namespace": tool_namespace} if tool_namespace else {} + return {**item_kwargs, "name": tool_name, **namespace_kwargs} + def _is_reasoning_end(self, chunk): delta: Final = chunk.choices[0].delta @@ -244,17 +253,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): else: fn_name = str(getattr(fn, "name", "") or "") fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", "")) - tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) output_index = self._get_or_assign_tool_output_index(call_id) if call_id not in self._tool_args_by_call_id: self._tool_args_by_call_id[call_id] = "" self._sequence_number += 1 - names = self._custom_tool_names - item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names) + item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, "", "in_progress") self._tool_item_id_by_call_id[call_id] = item_kwargs["id"] - if tool_namespace: - item_kwargs["namespace"] = tool_namespace event = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=output_index, @@ -315,7 +320,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): else: fn_name = str(getattr(fn, "name", "") or "") fn_args = serialize_tool_call_arguments(getattr(fn, "arguments", "")) - tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name) web_search_call = self._web_search_calls.get(call_id) if web_search_call is not None: if call_id not in self._queued_web_search_call_ids: @@ -330,11 +334,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if is_new_tool_call: self._tool_args_by_call_id[call_id] = "" self._sequence_number += 1 - names = self._custom_tool_names - item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names) + item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, "", "in_progress") self._tool_item_id_by_call_id[call_id] = item_kwargs["id"] - if tool_namespace: - item_kwargs["namespace"] = tool_namespace event = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=output_index, @@ -376,11 +377,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._pending_tool_events.append(done_event) self._sequence_number += 1 - names = self._custom_tool_names - item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, final_args, "completed", names) + item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, final_args, "completed") item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"]) - if tool_namespace: - item_kwargs["namespace"] = tool_namespace item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py index a7f65545eb2..bbd0cdf97e3 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py @@ -174,6 +174,21 @@ def test_custom_member_description_edit_lands_without_the_namespace_prefix_or_gr assert list(merged) == [{"type": "namespace", "name": "shell", "description": "Shell", "tools": [guarded_member]}] +def test_text_appended_after_the_grammar_block_lands_on_the_member_without_the_block(): + grammar = {"type": "grammar", "syntax": "lark", "definition": "start: X"} + custom_member = {"type": "custom", "name": "exec", "description": "Run a command", "format": grammar} + original = [{"type": "namespace", "name": "shell", "description": "Shell", "tools": [custom_member]}] + groups = _groups(original) + edited = copy.deepcopy(_flat(groups)) + edited[0]["function"]["description"] = edited[0]["function"]["description"] + " [checked]" + + merged = merge_guardrailed_tools(original, groups, edited) + + assert merged[0]["tools"][0]["description"] == "Run a command [checked]" + reflattened = _flat(_groups(merged)) + assert reflattened[0]["function"]["description"] == "Shell\n\nRun a command [checked]\n\nFormat:\n```lark\nstart: X\n```" + + def test_member_extras_edited_by_the_guardrail_land_on_that_member(): original = [{"type": "namespace", "name": "ns", "tools": [_function("read")]}] groups = _groups(original) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 1f497597a11..50a96c4271f 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -3877,6 +3877,83 @@ class TestEnsureOutputItemContentPartAdded: assert done.item.type == "custom_tool_call" assert done.item.input == "ls" + def test_streaming_namespaced_function_sharing_a_nested_custom_short_name_stays_a_function_call(self): + from litellm.responses.litellm_completion_transformation.custom_tools import extract_custom_tool_names + + iterator = self._make_iterator() + iterator.responses_api_request = { + "tools": [ + { + "type": "namespace", + "name": "alpha", + "tools": [ + { + "type": "custom", + "name": "run", + "format": {"type": "grammar", "syntax": "lark", "definition": "start: /.+/"}, + } + ], + }, + { + "type": "namespace", + "name": "beta", + "tools": [ + { + "type": "function", + "name": "run", + "parameters": {"type": "object", "properties": {"job_id": {"type": "string"}}}, + } + ], + }, + ] + } + iterator._custom_tool_names = extract_custom_tool_names(iterator.responses_api_request.get("tools")) + iterator._namespace_tool_names = LiteLLMCompletionResponsesConfig.namespace_tool_name_map( + iterator.responses_api_request.get("tools") + ) + function_call = {"id": "call_fn", "function": {"name": "beta__run", "arguments": '{"job_id":"42"}'}} + custom_call = {"id": "call_custom", "function": {"name": "run", "arguments": '{"content":"echo hi"}'}} + + iterator._queue_tool_call_delta_events([{"index": 0, **function_call}, {"index": 1, **custom_call}]) + iterator._queue_final_tool_call_done_events( + ModelResponse( + id="chatcmpl-run", + created=1, + model="us.openai.gpt-5.6", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id=call["id"], type="function", function=Function(**call["function"]) + ) + for call in (function_call, custom_call) + ], + ), + ) + ], + ) + ) + + items = [ + event.item + for event in iterator._pending_tool_events + if event.type in ("response.output_item.added", "response.output_item.done") + ] + function_items = [item for item in items if item.call_id == "call_fn"] + custom_items = [item for item in items if item.call_id == "call_custom"] + assert len(function_items) == 2 and len(custom_items) == 2 + assert all((item.type, item.name, item.namespace) == ("function_call", "run", "beta") for item in function_items) + assert function_items[-1].arguments == '{"job_id":"42"}' + assert all(item.type == "custom_tool_call" and item.name == "run" for item in custom_items) + assert all(getattr(item, "namespace", None) is None for item in custom_items) + assert custom_items[-1].input == "echo hi" + def test_streaming_unqualified_namespace_tool_calls_restore_namespace(self): """A unique nested tool name without the namespace still maps back.""" iterator = self._make_iterator() From 9fb94ea761f38feb350e5225e6b3467e3d641405 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:59:40 -0700 Subject: [PATCH 091/107] fix(exceptions): keep repeated litellm_proxy response headers on the rebuilt response httpx.Headers.items() comma-joins repeated header names, so the rebuilt response iterates multi_items() and keeps every value, matching what the raw openai client exposes on e.response.headers --- .../exception_mapping_utils.py | 3 ++- .../test_exception_mapping_utils.py | 17 ++++++++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index e8406b87777..70675966dfc 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -259,9 +259,10 @@ def _litellm_proxy_response( headers: Final = getattr(original_exception, "headers", None) if not isinstance(headers, Mapping) or not headers: return response + pairs: Final = headers.multi_items() if isinstance(headers, httpx.Headers) else headers.items() return httpx.Response( status_code=response.status_code, - headers={str(k): str(v) for k, v in headers.items()}, + headers=[(str(k), str(v)) for k, v in pairs], request=getattr(original_exception, "request", None), ) diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index a4869aac8fe..acc6248bf3e 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1373,7 +1373,7 @@ _GUARDRAIL_BLOCK_ERROR = { def _openai_handler_error( error_type: str, - headers: dict[str, str], + headers: dict[str, str] | list[tuple[str, str]], status_code: int = 400, message: str = _GUARDRAIL_BLOCK_ERROR["message"], ) -> OpenAIError: @@ -1435,3 +1435,18 @@ def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): assert exc_info.value.body["type"] == "vendor_specific_error" assert not exc_info.value.response.headers + + +def test_litellm_proxy_repeated_response_header_keeps_each_value(): + repeated = [("x-litellm-call-id", "call-guardrail"), ("set-cookie", "a=1"), ("set-cookie", "b=2")] + + with pytest.raises(litellm.BadRequestError) as exc_info: + exception_type( + model="claude-haiku-4-5", + original_exception=_openai_handler_error("None", repeated), + custom_llm_provider="litellm_proxy", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers.multi_items() == repeated From aaf924693a540572f90f512334ccbc056ce3554c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:13:50 -0700 Subject: [PATCH 092/107] fix(router): count num_retries_per_request across fallback hops num_retries_per_request has always capped the retries of one request with its fallback hops included. #40930 started reading the per-hop attempted_retries counter instead, and every fallback hop restarts that counter at zero, so a request could spend a fresh retry budget on each hop and the legacy fallback cap test started seeing the hop run. Router.log_retry now also keeps request_retry_count on the request metadata, incremented on every retry and fallback hop and never truncated the way previous_models is, and max_retries_per_request_hit reads that count. The flat retry records, the litellm_metadata coverage and caps above four from #40930 stay as they are, and the legacy test goes back to its previous_models == 0 assertion. --- litellm/__init__.py | 2 +- litellm/litellm_core_utils/core_helpers.py | 4 +- litellm/router.py | 6 ++- tests/local_testing/test_router_fallbacks.py | 17 ++---- .../test_router_helper_utils.py | 10 ++-- .../rust_bridge/test_lifecycle.py | 11 ++-- tests/test_litellm/test_router.py | 52 +++++++++++++++++++ tests/test_litellm/test_utils.py | 13 ++--- 8 files changed, 85 insertions(+), 30 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 261457d6889..ccfbf80369f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -538,7 +538,7 @@ context_window_fallbacks: Optional[List] = None content_policy_fallbacks: Optional[List] = None allowed_fails: int = 3 allow_dynamic_callback_disabling: bool = True -num_retries_per_request: Optional[int] = None # cap on Router retries of one model group; resets per fallback hop +num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries) ####### SECRET MANAGERS ##################### secret_manager_client: Optional[Any] = ( None # list of instantiated key management clients - e.g. azure kv, infisical, etc. diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 6e76bf9d49e..15380bc5d57 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -309,8 +309,8 @@ def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_re metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) if not isinstance(metadata, Mapping): return False - attempted_retries: Final = metadata.get("attempted_retries") - return type(attempted_retries) is int and 0 < attempted_retries and num_retries_per_request <= attempted_retries + retry_count: Final = metadata.get("request_retry_count") + return type(retry_count) is int and 0 < retry_count and num_retries_per_request <= retry_count def get_or_create_metadata_bucket( diff --git a/litellm/router.py b/litellm/router.py index 1665583386f..ff50cac1328 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8378,7 +8378,8 @@ class Router: def log_retry(self, kwargs: dict, e: Exception) -> dict: """ - When a retry or fallback happens, record which model group, deployment and attempt just failed and why + When a retry or fallback happens, record which model group, deployment and attempt just failed and why, + and count it toward the request-wide num_retries_per_request cap """ from litellm.types.router import RetryAttemptRecord @@ -8402,7 +8403,10 @@ class Router: else () ) breadcrumbs: Final = (*kept_breadcrumbs, attempt_record) + earlier_retry_count: Final = request_metadata.get("request_retry_count") + request_retry_count: Final = (earlier_retry_count if type(earlier_retry_count) is int else 0) + 1 kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict + kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap return kwargs def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int: diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index 5d7955ad34a..82b832f89fd 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -5,7 +5,6 @@ import asyncio import os import time import traceback -from typing import Final import pytest @@ -14,7 +13,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger -from litellm.types.router import RetryAttemptRecord from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE @@ -23,15 +21,15 @@ class MyCustomHandler(CustomLogger): success: bool = False failure: bool = False previous_models: int = 0 - previous_model_records: tuple[RetryAttemptRecord, ...] = () def log_pre_api_call(self, model, messages, kwargs): print(f"Pre-API Call") print( f"previous_models: {kwargs['litellm_params']['metadata'].get('previous_models', None)}" ) - self.previous_model_records = tuple(kwargs["litellm_params"]["metadata"].get("previous_models", ())) - self.previous_models = len(self.previous_model_records) + self.previous_models = len( + kwargs["litellm_params"]["metadata"].get("previous_models", []) + ) # {"previous_models": [{"model": litellm_model_name, "exception_type": AuthenticationError, "exception_string": }]} print(f"self.previous_models: {self.previous_models}") def log_post_api_call(self, kwargs, response_obj, start_time, end_time): @@ -720,14 +718,7 @@ async def test_async_fallbacks_max_retries_per_request(): await asyncio.sleep( 0.05 ) # allow a delay as success_callbacks are on a separate thread - records: Final = customHandler.previous_model_records - assert customHandler.previous_models == len(records) - assert records - assert {record["model_group"] for record in records} == {"azure/gpt-3.5-turbo"} - assert next(record["exception_type"] for record in records if record["attempted_retries"] == 0) == "AuthenticationError" - refused_retries: Final = tuple(record for record in records if record["attempted_retries"]) - assert refused_retries - assert all("Max retries per request hit!" in record["exception_string"] for record in refused_retries) + assert customHandler.previous_models == 0 # 0 retries, 0 fallback router.reset() except litellm.Timeout as e: pass diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 5b06c5fdb01..c7e577366c3 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -631,9 +631,11 @@ def test_deployment_callback_respects_cooldown_time(model_list): @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) def test_log_retry(model_list, metadata_key): - """log_retry appends one flat record per failed attempt and copies neither the request kwargs nor - the request metadata into it""" + """log_retry appends one flat record per failed attempt, copies neither the request kwargs nor the + request metadata into it, and counts every failed attempt of the request independently of the + per-hop attempted_retries""" router = Router(model_list=model_list) + rate_limit_error = litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo") new_kwargs = router.log_retry( kwargs={ "model": "gpt-3.5-turbo", @@ -641,7 +643,7 @@ def test_log_retry(model_list, metadata_key): "messages": [{"role": "user", "content": "hi"}], metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"}, }, - e=litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo"), + e=rate_limit_error, ) assert json.loads(json.dumps(new_kwargs[metadata_key]["previous_models"])) == [ { @@ -652,6 +654,8 @@ def test_log_retry(model_list, metadata_key): "attempted_retries": 2, } ] + assert new_kwargs[metadata_key]["request_retry_count"] == 1 + assert router.log_retry(kwargs=new_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 2 def test_update_usage(model_list): diff --git a/tests/test_litellm/rust_bridge/test_lifecycle.py b/tests/test_litellm/rust_bridge/test_lifecycle.py index 1f0b5591c2b..d73385621d5 100644 --- a/tests/test_litellm/rust_bridge/test_lifecycle.py +++ b/tests/test_litellm/rust_bridge/test_lifecycle.py @@ -8,7 +8,7 @@ from litellm.rust_bridge.lifecycle import check_limits @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize( - "cap, attempted_retries, refused", + "cap, request_retry_count, refused", [(5, 5, True), (5, 4, False), (0, 0, False), (0, 1, True)], ids=[ "cap-above-four-reached", @@ -17,12 +17,15 @@ from litellm.rust_bridge.lifecycle import check_limits "cap-of-zero-refuses-first-retry", ], ) -def test_check_limits_reads_attempted_retries( - monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, attempted_retries: int, refused: bool +def test_check_limits_reads_request_retry_count( + monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, request_retry_count: int, refused: bool ) -> None: monkeypatch.setattr(litellm, "num_retries_per_request", cap) monkeypatch.setattr(litellm, "max_budget", None) - kwargs: Final = {"model": "mistral/mistral-ocr-latest", metadata_key: {"attempted_retries": attempted_retries}} + kwargs: Final = { + "model": "mistral/mistral-ocr-latest", + metadata_key: {"request_retry_count": request_retry_count}, + } if refused: with pytest.raises(RuntimeError, match="Max retries per request hit!"): check_limits(kwargs) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 3cabcd71627..577332727d7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11066,6 +11066,58 @@ async def test_num_retries_per_request_stops_retries_at_caps_above_four(monkeypa ] +def _failing_group_with_healthy_fallback_router(num_retries): + return litellm.Router( + model_list=[ + { + "model_name": "broken-group", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + "mock_response": "litellm.InternalServerError", + }, + }, + { + "model_name": "healthy-group", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-fake", "mock_response": "ok"}, + }, + ], + fallbacks=[{"broken-group": ["healthy-group"]}], + num_retries=num_retries, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cap, hop_refused", [(2, True), (4, False)], ids=["cap-spent-before-the-hop", "cap-not-reached-by-the-hop"] +) +async def test_num_retries_per_request_counts_retries_across_fallback_hops(monkeypatch, cap, hop_refused): + """num_retries_per_request caps the retries of one request, fallback hops included. Each hop starts a + fresh per-hop attempted_retries at zero, so a cap read from that counter let every hop retry from zero + and a request could spend far more retries than the cap allows.""" + monkeypatch.setattr(litellm, "num_retries_per_request", cap) + router = _failing_group_with_healthy_fallback_router(num_retries=1) + recorder = _FallbackAttemptRecorder() + litellm.callbacks.append(recorder) + try: + request = router.acompletion(model="broken-group", messages=[{"role": "user", "content": "hi"}]) + if not hop_refused: + assert (await request).choices[0].message.content == "ok" + return + with pytest.raises(litellm.InternalServerError): + await request + finally: + litellm.callbacks.remove(recorder) + + assert recorder.failed_targets == ["healthy-group"] + hop_refusals = [ + record["attempted_retries"] + for record in recorder.breadcrumbs_per_target[0] + if record["model_group"] == "healthy-group" and "Max retries per request hit!" in record["exception_string"] + ] + assert hop_refusals == [0, 1] + + @pytest.mark.asyncio async def test_fallback_traceback_stays_available_at_debug_level(): """Dropping the stack from the ERROR line is only safe because the fallback path still diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d5feda6f892..3d60fe86d9c 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4077,10 +4077,11 @@ class TestMetadataNoneHandling: _RETRY_CAP_CASES: Final = ( - pytest.param(5, {"attempted_retries": 5}, True, id="cap-above-four-reached"), - pytest.param(5, {"attempted_retries": 4}, False, id="cap-above-four-not-reached"), - pytest.param(0, {"attempted_retries": 0}, False, id="first-attempt-passes-cap-of-zero"), - pytest.param(0, {"attempted_retries": 1}, True, id="cap-of-zero-refuses-first-retry"), + pytest.param(5, {"request_retry_count": 5}, True, id="cap-above-four-reached"), + pytest.param(5, {"request_retry_count": 4}, False, id="cap-above-four-not-reached"), + pytest.param(0, {"request_retry_count": 0}, False, id="first-attempt-passes-cap-of-zero"), + pytest.param(0, {"request_retry_count": 1}, True, id="cap-of-zero-refuses-first-retry"), + pytest.param(0, {"attempted_retries": 1}, False, id="per-hop-attempted-retries-is-not-the-cap"), pytest.param(5, {"previous_models": ("a", "b", "c", "d", "e")}, False, id="breadcrumb-count-is-not-the-cap"), pytest.param(5, None, False, id="metadata-none"), ) @@ -4098,7 +4099,7 @@ def _capped_completion_kwargs(metadata_key: str, metadata: object) -> dict[str, @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) -def test_num_retries_per_request_reads_attempted_retries_sync(monkeypatch, metadata_key, cap, metadata, refused): +def test_num_retries_per_request_reads_request_retry_count_sync(monkeypatch, metadata_key, cap, metadata, refused): monkeypatch.setattr(litellm, "num_retries_per_request", cap) kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) if refused: @@ -4111,7 +4112,7 @@ def test_num_retries_per_request_reads_attempted_retries_sync(monkeypatch, metad @pytest.mark.asyncio @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) @pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES) -async def test_num_retries_per_request_reads_attempted_retries_async(monkeypatch, metadata_key, cap, metadata, refused): +async def test_num_retries_per_request_reads_request_retry_count_async(monkeypatch, metadata_key, cap, metadata, refused): monkeypatch.setattr(litellm, "num_retries_per_request", cap) kwargs: Final = _capped_completion_kwargs(metadata_key, metadata) if refused: From 2bf44ed35480b17d7757d8817985b73147a486f0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:18:47 -0700 Subject: [PATCH 093/107] fix(guardrails): reject tool_use rewrites that are not JSON objects --- .../chat/guardrail_translation/handler.py | 36 ++++++++++++------- .../test_anthropic_guardrail_handler.py | 14 +++++--- 2 files changed, 33 insertions(+), 17 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9f3c85b555a..b222548f4ec 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -234,19 +234,20 @@ def _write_back_message_text(message: _WritableMessage, target: MessageTextTarge _TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object]) -def _write_back_tool_use(message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape) -> None: +def _rewritten_tool_use_input(arguments: str) -> Mapping[str, object] | None: + try: + return _TOOL_USE_INPUT_ADAPTER.validate_json(arguments) + except ValidationError: + return None + + +def _write_back_tool_use( + message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape, rewritten_input: Mapping[str, object] +) -> None: content: Final = message.get("content", None) block: Final = content[target.content_idx] if isinstance(content, list) else None if not isinstance(block, dict): return - try: - rewritten_input: Final = _TOOL_USE_INPUT_ADAPTER.validate_json(shape.arguments) - except ValidationError: - verbose_proxy_logger.warning( - "Anthropic Messages: guardrail returned arguments that are not a JSON object for tool_use %s; keeping its input", - block.get("id"), - ) - return block["input"] = rewritten_input # mutable-ok: guardrails rewrite the caller's request payload in place if shape.name is not None and shape.name != block.get("name"): block["name"] = shape.name # mutable-ok: guardrails rewrite the caller's request payload in place @@ -688,6 +689,7 @@ class AnthropicMessagesHandler(BaseTranslation): scanned_tool_calls=scanned_tool_calls, pre_guardrail_tool_calls=pre_guardrail_tool_calls, returned_tool_calls=guardrailed_inputs.get("tool_calls"), + guardrail_name=guardrail_to_apply.guardrail_name, ) verbose_proxy_logger.debug("Anthropic Messages: Processed input messages: %s", messages) @@ -1116,15 +1118,25 @@ class AnthropicMessagesHandler(BaseTranslation): scanned_tool_calls: tuple[ScannedToolCall, ...], pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], returned_tool_calls: Sequence[object] | None, + guardrail_name: str | None, ) -> None: post_guardrail_tool_calls: Final = _tool_call_shapes( returned_tool_calls if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls) else tuple(item.tool_call for item in scanned_tool_calls) ) - for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls): - if before != after: - _write_back_tool_use(messages[item.target.msg_idx], item.target, after) + rewritten: Final = tuple( + (item, after, _rewritten_tool_use_input(after.arguments)) + for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if before != after + ) + applicable: Final = tuple( + (item, after, rewritten_input) for item, after, rewritten_input in rewritten if rewritten_input is not None + ) + if len(applicable) != len(rewritten): + raise unappliable_request_rewrite(guardrail_name) + for item, after, rewritten_input in applicable: + _write_back_tool_use(messages[item.target.msg_idx], item.target, after, rewritten_input) async def process_output_response( self, diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 287674afaa7..b73ef6453fa 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2328,16 +2328,20 @@ class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: assert data["messages"][2]["content"][0]["tool_use_id"] == "toolu_01" @pytest.mark.asyncio - async def test_non_json_rewritten_arguments_keep_the_tool_use_input(self): + async def test_non_json_rewritten_arguments_are_rejected_by_name(self): + from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite + handler = AnthropicMessagesHandler() guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]") data = self._tool_use_conversation(system="You are a careful agent harness.") + original = json.loads(json.dumps(data)) - await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + with pytest.raises(UnappliableRequestRewrite) as excinfo: + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - assert data["messages"][1]["content"][0]["input"] == { - "cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity" - } + assert excinfo.value.guardrail_name == "scan-only-capture" + assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched" + assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched" @pytest.mark.asyncio async def test_scan_only_tool_results_keeps_system_and_tool_use_out(self): From 398300c4e70430a2ef1323345efa569a83e2a7de Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 14 Sep 2026 02:22:55 -0700 Subject: [PATCH 094/107] fix(router): honor team and key provider weights --- litellm/proxy/_types.py | 12 +- litellm/proxy/auth/login_utils.py | 2 + litellm/proxy/auth/user_api_key_auth.py | 1 + litellm/proxy/common_request_processing.py | 9 ++ litellm/proxy/litellm_pre_call_utils.py | 2 + .../internal_user_endpoints.py | 2 +- .../key_management_endpoints.py | 52 ++++++- .../management_endpoints/router_weights.py | 129 ++++++++++++++++++ .../management_endpoints/team_endpoints.py | 16 +++ litellm/proxy/management_endpoints/ui_sso.py | 1 + litellm/proxy/proxy_server.py | 2 + litellm/router.py | 26 ++-- litellm/router_strategy/simple_shuffle.py | 118 ++++++++-------- litellm/types/router.py | 2 + litellm/types/router_weights.py | 30 ++++ litellm/types/utils.py | 1 + .../management_endpoints/test_common_utils.py | 42 ++++++ .../test_key_management_endpoints.py | 54 +++++++- .../test_team_endpoints.py | 20 +++ .../proxy/test_common_request_processing.py | 39 ++++++ .../proxy/test_litellm_pre_call_utils.py | 6 + tests/test_litellm/proxy/test_proxy_types.py | 10 ++ .../router_strategy/test_simple_shuffle.py | 50 +++++++ tests/test_litellm/test_utils.py | 10 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 18 +-- 25 files changed, 561 insertions(+), 93 deletions(-) create mode 100644 litellm/proxy/management_endpoints/router_weights.py create mode 100644 litellm/types/router_weights.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ffc41a9d7ae..d40f518d4b8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4,11 +4,12 @@ import os from collections.abc import Callable, Mapping from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple +from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple import httpx from pydantic import ( BaseModel, + BeforeValidator, ConfigDict, Field, Json, @@ -47,6 +48,7 @@ from litellm.types.proxy.carried_budget_state import ( ) from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.types.router import RouterErrors, UpdateRouterConfig +from litellm.types.router_weights import validate_router_settings_dict from litellm.types.secret_managers.main import KeyManagementSystem from litellm.types.utils import ( CallTypes, @@ -1981,8 +1983,14 @@ class OrgMember(MemberBase): from litellm.models.team import TeamBase as TeamBase # noqa: E402 +RouterSettingsDict = Annotated[ + dict[str, object], + BeforeValidator(validate_router_settings_dict, json_schema_input_type=UpdateRouterConfig), +] + class NewTeamRequest(TeamBase): + router_settings: RouterSettingsDict | None = None model_aliases: dict | None = None tags: list | None = None guardrails: list[str] | None = None @@ -2080,7 +2088,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None enforced_batch_output_expires_after: dict | None = None enforced_file_expires_after: dict | None = None - router_settings: dict | None = None + router_settings: RouterSettingsDict | None = None access_group_ids: list[str] | None = None budget_limits: list[BudgetLimitEntry] | None = None # multiple concurrent budget windows default_team_member_models: list[str] | None = None # default allowed_models seeded onto new team members diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index c0a76a4fc20..b7064802878 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -249,6 +249,7 @@ async def authenticate_user( if os.getenv("DATABASE_URL") is not None: response = await generate_key_helper_fn( + llm_router=None, request_type="key", **{ "user_role": LitellmUserRoles.PROXY_ADMIN, @@ -324,6 +325,7 @@ async def authenticate_user( await _rehash_password_if_needed(_user_row.user_id, password, _password) if os.getenv("DATABASE_URL") is not None: response = await generate_key_helper_fn( + llm_router=None, request_type="key", **{ "user_role": user_role, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 25570ab220a..9491f77ecfc 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -878,6 +878,7 @@ async def _auto_register_jwt_mapping( # the NOT NULL @id constraint. Every successful key-creation caller (e.g. # /key/generate) passes table_name="key" explicitly. key_data: Final = await generate_key_helper_fn( + llm_router=None, request_type="key", table_name="key", team_id=team_id, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 0ad86479aac..4a4daa68cce 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -14,6 +14,7 @@ import httpx import orjson from fastapi import HTTPException, Request, status from fastapi.responses import JSONResponse, Response, StreamingResponse +from pydantic import ValidationError from starlette.types import Receive, Scope, Send import litellm @@ -76,6 +77,7 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError +from litellm.types.router_weights import validate_router_weights _LateResponseT = TypeVar("_LateResponseT", bound=Response) _LlmCallT = TypeVar("_LlmCallT") @@ -1939,6 +1941,13 @@ class ProxyBaseLLMRequestProcessing: # This avoids expensive Router instantiation on each request if router_settings is not None: self.data["router_settings_override"] = router_settings + try: + self.data["_router_weights"] = validate_router_weights(router_settings.get("weights")) + except ValidationError: + self.data["_router_weights"] = None + verbose_proxy_logger.warning( + "Ignoring invalid saved router weights; update team/key router_settings" + ) alias_target: Final = await _resolve_per_request_model_group_alias( requested_model=self.data.get("model"), router_settings=router_settings, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 59971e54e46..ea09ab734a7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -221,6 +221,8 @@ LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset( ) _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( + "weights", + "_router_weights", "proxy_server_request", "standard_logging_object", "secret_fields", diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index e3efda507f6..ba7a3309a90 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -567,7 +567,7 @@ async def new_user( teams = check_if_default_team_set() organization_ids: Final = cast(list[str] | None, data_json.pop("organizations", None)) - response: Final = await generate_key_helper_fn(request_type="user", **data_json) + response: Final = await generate_key_helper_fn(request_type="user", **data_json, llm_router=None) # Admin UI Logic # Add User to Team and Organization # if team_id passed add this user to the team diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 95ccb7bbe0b..339933807f1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -96,6 +96,7 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) +from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights from litellm.proxy.management_helpers.access_group_key_sync import ( sync_key_access_group_membership, sync_key_regeneration_access_group_membership, @@ -201,6 +202,10 @@ class _KeyUpdateResult(TypedDict): data: ReadOnly[Mapping[str, object]] +class _StoredKeyRouterSettings(BaseModel): + router_settings: Mapping[str, object] | None = None + + class _KeyRowWhere(TypedDict): token: ReadOnly[str] @@ -1330,7 +1335,7 @@ async def _common_key_generation_helper( prisma_client=prisma_client, ) - response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key") + response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key", llm_router=llm_router) response["soft_budget"] = data.soft_budget # include the user-input soft budget in the response @@ -2234,7 +2239,26 @@ async def _update_key_row_with_soft_budget( async def prepare_key_update_data( data: UpdateKeyRequest | RegenerateKeyRequest, existing_key_row: LiteLLM_VerificationToken, + *, + prisma_client: PrismaClient | None = None, + llm_router: Router | None = None, ): + if data.router_settings is not None or ( + "router_settings" not in data.model_fields_set + and "team_id" in data.model_fields_set + and data.team_id != existing_key_row.team_id + ): + effective_settings: Final = ( + data.router_settings + if data.router_settings is not None + else _StoredKeyRouterSettings.model_validate(existing_key_row, from_attributes=True).router_settings + ) + await validate_router_settings_weights( + effective_settings, + team_id=data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id, + prisma_client=prisma_client, + llm_router=llm_router, + ) data_json: Final[dict] = data.model_dump(exclude_unset=True) data_json.pop("key", None) data_json.pop("new_key", None) @@ -2575,7 +2599,9 @@ async def _process_single_key_update( ) # Prepare update data - non_default_values = await prepare_key_update_data(data=update_key_request, existing_key_row=existing_key_row) + non_default_values = await prepare_key_update_data( + data=update_key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router + ) # Update key in database if prisma_client is None: @@ -3093,7 +3119,9 @@ async def update_key_fn( # Enforce upperbound key params on update (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) - non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row) + non_default_values: Final = await prepare_key_update_data( + data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router + ) # Only validate key_alias format if it's actually being changed new_key_alias: Final = non_default_values.get("key_alias", None) @@ -4137,15 +4165,24 @@ async def generate_key_helper_fn( object_permission: LiteLLM_ObjectPermissionBase | None = None, auto_rotate: bool | None = None, rotation_interval: str | None = None, - router_settings: dict | None = None, + router_settings: dict[str, object] | None = None, access_group_ids: list[str] | None = None, budget_limits: list | None = None, # multiple concurrent budget windows + *, + llm_router: Router | None = None, ): from litellm.proxy.proxy_server import premium_user, prisma_client if prisma_client is None: raise Exception("Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys ") + await validate_router_settings_weights( + router_settings, + team_id=team_id, + prisma_client=prisma_client, + llm_router=llm_router, + ) + if token is None: if key is not None: token = key @@ -5070,6 +5107,7 @@ async def _insert_deprecated_key( async def _execute_virtual_key_regeneration( *, prisma_client: PrismaClient, + llm_router: Router | None = None, key_in_db: LiteLLM_VerificationToken, hashed_api_key: str, key: str, @@ -5129,7 +5167,9 @@ async def _execute_virtual_key_regeneration( if data is not None: # Enforce upperbound key params on regenerate (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) - non_default_values = await prepare_key_update_data(data=data, existing_key_row=key_in_db) + non_default_values = await prepare_key_update_data( + data=data, existing_key_row=key_in_db, prisma_client=prisma_client, llm_router=llm_router + ) # Only validate key_alias format if it's actually being changed new_key_alias: Final = non_default_values.get("key_alias") if new_key_alias != key_in_db.key_alias: @@ -5268,6 +5308,7 @@ async def regenerate_key_fn( try: from litellm.proxy.proxy_server import ( hash_token, + llm_router, master_key, premium_user, prisma_client, @@ -5456,6 +5497,7 @@ async def regenerate_key_fn( return await _execute_virtual_key_regeneration( prisma_client=prisma_client, + llm_router=llm_router, key_in_db=_key_in_db, hashed_api_key=hashed_api_key, key=key, diff --git a/litellm/proxy/management_endpoints/router_weights.py b/litellm/proxy/management_endpoints/router_weights.py new file mode 100644 index 00000000000..99b368808c3 --- /dev/null +++ b/litellm/proxy/management_endpoints/router_weights.py @@ -0,0 +1,129 @@ +from abc import abstractmethod +from collections.abc import Mapping +from typing import Annotated, Final, Protocol + +from fastapi import HTTPException +from pydantic import BaseModel, BeforeValidator, ValidationError + +from litellm.repositories.prisma_protocols import TableActions +from litellm.types.router_weights import RouterWeights + + +class _StoredModel(Protocol): + @property + @abstractmethod + def model_id(self) -> str: + pass + + +class _ModelDb(Protocol): + @property + @abstractmethod + def litellm_proxymodeltable(self) -> TableActions[_StoredModel]: + pass + + +class _PrismaClient(Protocol): + @property + @abstractmethod + def db(self) -> _ModelDb: + pass + + +class _Router(Protocol): + @abstractmethod + def get_deployment(self, model_id: str) -> object | None: + pass + + +class _RouterWeightSettings(BaseModel): + weights: RouterWeights | None = None + + +class _RouterWeightModelInfo(BaseModel): + team_id: str | None = None + db_model: bool | None = None + team_public_model_name: str | None = None + + +def _router_weight_model_info(value: object) -> _RouterWeightModelInfo: + if isinstance(value, str): + return _RouterWeightModelInfo.model_validate_json(value) + return _RouterWeightModelInfo.model_validate(value or {}, from_attributes=True) + + +class _RouterWeightDeployment(BaseModel): + model_name: str + model_info: Annotated[_RouterWeightModelInfo, BeforeValidator(_router_weight_model_info)] + + +def _validate_router_weight_reference( + model_group: str, + deployment_id: str, + team_id: str | None, + stored: _RouterWeightDeployment | None, + configured: object | None, +) -> None: + reference: Final = ( + stored + if stored is not None + else ( + _RouterWeightDeployment.model_validate(configured, from_attributes=True) if configured is not None else None + ) + ) + if ( + reference is None + or (stored is None and reference.model_info.db_model) + or (reference.model_info.team_id is not None and reference.model_info.team_id != team_id) + ): + raise HTTPException(status_code=400, detail=f"Unknown deployment ID in router weights: {deployment_id}") + canonical_group: Final = ( + reference.model_info.team_public_model_name if reference.model_info.team_id is not None else None + ) or reference.model_name + if model_group != canonical_group: + raise HTTPException( + status_code=400, + detail=f"Deployment {deployment_id} does not belong to model group {model_group}", + ) + + +async def validate_router_settings_weights( + router_settings: BaseModel | Mapping[str, object] | None, + *, + team_id: str | None, + prisma_client: _PrismaClient | None, + llm_router: _Router | None, +) -> None: + try: + weights: Final = ( + _RouterWeightSettings.model_validate(router_settings, from_attributes=True).weights + if router_settings is not None + else None + ) + except ValidationError: + raise HTTPException( + status_code=400, + detail="Invalid router weights. Replace or clear router_settings.weights.", + ) from None + if not weights: + return + deployment_ids: Final = frozenset(deployment_id for group in weights.values() for deployment_id in group) + if not deployment_ids: + return + if prisma_client is None: + raise HTTPException(status_code=503, detail="Database unavailable while validating router weights") + stored_models: Final = await prisma_client.db.litellm_proxymodeltable.find_many( + where={"model_id": {"in": list(deployment_ids)}} + ) + stored_by_id: Final = { + row.model_id: _RouterWeightDeployment.model_validate(row, from_attributes=True) for row in stored_models + } + for model_group, group_weights in weights.items(): + for deployment_id in group_weights: + _validate_router_weight_reference( + model_group, + deployment_id, + team_id, + stored_by_id.get(deployment_id), + llm_router.get_deployment(model_id=deployment_id) if llm_router is not None else None, + ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 9350d2cd691..2d04a4d1e04 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -112,6 +112,7 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.organization_endpoints import ( add_member_to_organization, ) +from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, ) @@ -1288,6 +1289,7 @@ async def new_team( create_audit_log_for_update, general_settings, litellm_proxy_admin_name, + llm_router, prisma_client, user_api_key_cache, ) @@ -1462,6 +1464,13 @@ async def new_team( user_api_key_dict=user_api_key_dict, ) + await validate_router_settings_weights( + data.router_settings, + team_id=data.team_id, + prisma_client=prisma_client, + llm_router=llm_router, + ) + ## ADD TO MODEL TABLE _model_id = None if data.model_aliases is not None and isinstance(data.model_aliases, dict): @@ -2075,6 +2084,13 @@ async def update_team( user_api_key_dict=user_api_key_dict, ) + await validate_router_settings_weights( + data.router_settings, + team_id=data.team_id, + prisma_client=prisma_client, + llm_router=llm_router, + ) + _existing_team_metadata: Final[object] = getattr(existing_team_row, "metadata", None) enforce_output_token_estimates_are_admin_only( data=data, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 091dccf1433..329443148a2 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -3592,6 +3592,7 @@ class SSOAuthenticationHandler: verbose_proxy_logger.info("user_defined_values for creating ui key: %s", user_defined_values) response: Final = await generate_key_helper_fn( + llm_router=None, request_type="key", duration=LITELLM_UI_SESSION_DURATION, key_max_budget=litellm.max_ui_session_budget, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1c931863a2f..f81a3166a28 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9546,6 +9546,7 @@ class ProxyStartupEvent: gate the first duration window. """ await generate_key_helper_fn( + llm_router=llm_router, request_type="user", table_name="user", user_id=LITELLM_PROXY_BUDGET_NAME, @@ -16290,6 +16291,7 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str: global master_key, general_settings response: Final = await generate_key_helper_fn( + llm_router=llm_router, request_type="key", **{ "user_role": user_obj.user_role, diff --git a/litellm/router.py b/litellm/router.py index 1665583386f..6751711d690 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4849,6 +4849,7 @@ class Router: model=model, messages=messages, specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, ) data: Final = deployment["litellm_params"].copy() @@ -5163,13 +5164,11 @@ class Router: return healthy_deployments[0] # Use simple_shuffle for weighted selection - return cast( - GuardrailTypedDict, - simple_shuffle( - llm_router_instance=self, - healthy_deployments=healthy_deployments, - model=guardrail_name, - ), + return simple_shuffle( + resolve_model_alias=self._get_model_from_alias, + healthy_deployments=healthy_deployments, + model=guardrail_name, + request_kwargs=None, ) async def _ageneric_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs): @@ -13045,9 +13044,10 @@ class Router: start_time: Final = time.time() if strategy == "simple-shuffle": return simple_shuffle( - llm_router_instance=self, + resolve_model_alias=self._get_model_from_alias, healthy_deployments=healthy_deployments, model=model, + request_kwargs=request_kwargs, ) deployment: Final = await self._select_deployment_async( strategy=strategy, @@ -13190,9 +13190,10 @@ class Router: start_time: Final = time.perf_counter() if strategy == "simple-shuffle": return simple_shuffle( - llm_router_instance=self, + resolve_model_alias=self._get_model_from_alias, healthy_deployments=pass_through_deployments, model=model, + request_kwargs=request_kwargs, ) deployment: Final = await self._select_deployment_async( strategy=strategy, @@ -13888,9 +13889,10 @@ class Router: # if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm ############## Check 'weight' param set for weighted pick ################# return simple_shuffle( - llm_router_instance=self, + resolve_model_alias=self._get_model_from_alias, healthy_deployments=healthy_deployments, model=model, + request_kwargs=request_kwargs, ) deployment: Final = self._select_deployment_sync( strategy=strategy, @@ -13958,6 +13960,7 @@ class Router: messages=messages, input=input, specific_deployment=specific_deployment, + request_kwargs=request_kwargs, ) strategy, strategy_selector = self._get_routing_context(model, request_kwargs) @@ -14040,9 +14043,10 @@ class Router: # 6. Apply load balancing strategy if strategy == "simple-shuffle": return simple_shuffle( - llm_router_instance=self, + resolve_model_alias=self._get_model_from_alias, healthy_deployments=pass_through_deployments, model=model, + request_kwargs=request_kwargs, ) deployment: Final = self._select_deployment_sync( strategy=strategy, diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index 860e89cea22..4f2c5e8d933 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -1,71 +1,67 @@ -""" -Returns a random deployment from the list of healthy deployments. +"""Choose among eligible deployments using request weights, then global metrics.""" -If weights are provided, it will return a deployment based on the weights. - -""" +from __future__ import annotations +import logging import random -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Callable, Mapping, Sequence +from itertools import chain +from typing import Final, TypeVar -from litellm._logging import verbose_router_logger +from litellm.types.router_weights import validate_router_weights -if TYPE_CHECKING: - from litellm.router import Router as _Router +_DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) +_ROUTER_LOGGER: Final = logging.getLogger("LiteLLM Router") - LitellmRouter = _Router -else: - LitellmRouter = Any + +def _metric_weight(deployment: Mapping[str, object], metric: str) -> float: + params: Final = deployment.get("litellm_params") + value: Final = params.get(metric) if isinstance(params, Mapping) else None + if value is None: + return 0.0 + if isinstance(value, (int, float)): + return float(value) + raise TypeError(f"Deployment {metric} must be numeric") + + +def _scoped_weights( + deployments: Sequence[Mapping[str, object]], + model: str, + request_kwargs: Mapping[str, object] | None, +) -> tuple[float, ...]: + settings: Final = validate_router_weights((request_kwargs or {}).get("_router_weights")) + model_weights: Final = settings.get(model) if settings is not None else None + if not model_weights: + return () + return tuple( + model_weights.get(str(info.get("id")), 0.0) if isinstance(info, Mapping) else 0.0 + for deployment in deployments + for info in (deployment.get("model_info"),) + ) def simple_shuffle( - llm_router_instance: LitellmRouter, - healthy_deployments: list[Any] | dict[Any, Any], + resolve_model_alias: Callable[[str], str | None], + healthy_deployments: Sequence[_DeploymentT], model: str, -) -> dict: - """ - Returns a random deployment from the list of healthy deployments. - - If weights are provided, it will return a deployment based on the weights. - - If users pass `rpm` or `tpm`, we do a random weighted pick - based on `rpm`/`tpm`. - - Args: - llm_router_instance: LitellmRouter instance - healthy_deployments: List of healthy deployments - model: Model name - - Returns: - Dict: A single healthy deployment - """ - - ############## Check if 'weight' or 'rpm' or 'tpm' param set for a weighted pick ################# - for weight_by in ["weight", "rpm", "tpm"]: - if any(m["litellm_params"].get(weight_by) is not None for m in healthy_deployments): - weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments] - verbose_router_logger.debug("\nweight %s", weights) - total_weight = sum(weights) - if total_weight <= 0: - # All remaining candidates have weight 0 for this metric (e.g. - # after a weighted-failover exclusion left only zero-weight - # backups). Skip to the next metric (rpm/tpm) which may still - # provide a meaningful weighted pick; if none do, we fall - # through to the uniform random pick at the end. - continue - weights = [weight / total_weight for weight in weights] - verbose_router_logger.debug("\n weights %s by %s", weights, weight_by) - # Perform weighted random pick - selected_index = random.choices(range(len(weights)), weights=weights)[0] - verbose_router_logger.debug("\n selected index, %s", selected_index) - deployment = healthy_deployments[selected_index] - verbose_router_logger.info( - "get_available_deployment for model: %s, Selected deployment: %s for model: %s", - model, - llm_router_instance.print_deployment(deployment) or deployment[0], - model, - ) - return deployment or deployment[0] - - ############## No RPM/TPM passed, we do a random pick ################# - item: Final = random.choice(healthy_deployments) - return item or item[0] + request_kwargs: Mapping[str, object] | None, +) -> _DeploymentT: + resolved_model: Final = resolve_model_alias(model) or model + weight_sets: Final = chain( + (_scoped_weights(healthy_deployments, resolved_model, request_kwargs),), + ( + tuple(_metric_weight(deployment, metric) for deployment in healthy_deployments) + for metric in ("weight", "rpm", "tpm") + ), + ) + for weights in weight_sets: + largest = max(weights, default=0.0) + if largest <= 0: + continue + normalized = tuple(weight / largest for weight in weights) + if sum(normalized) <= 0: + continue + selected = random.choices(healthy_deployments, weights=normalized)[0] + _ROUTER_LOGGER.info("Selected deployment for model %s: %s", model, selected.get("model_info")) + return selected + return random.choice(healthy_deployments) diff --git a/litellm/types/router.py b/litellm/types/router.py index 0aefc07ae4b..49109e4fbbe 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -15,6 +15,7 @@ from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_c from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.types.router_weights import RouterWeights if TYPE_CHECKING: from litellm.router import Router @@ -146,6 +147,7 @@ class UpdateRouterConfig(BaseModel): context_window_fallbacks: list[dict] | None = None model_group_alias: dict[str, str | dict] | None = {} enable_tag_filtering: bool | None = None + weights: RouterWeights | None = None tag_routing_prefix: str | None = None optional_pre_call_checks: OptionalPreCallChecks | None = None diff --git a/litellm/types/router_weights.py b/litellm/types/router_weights.py new file mode 100644 index 00000000000..fa156661564 --- /dev/null +++ b/litellm/types/router_weights.py @@ -0,0 +1,30 @@ +from collections.abc import Mapping +from typing import Annotated, Final + +from pydantic import AfterValidator, Field, TypeAdapter + + +def _validate_positive_router_weights(weights: Mapping[str, Mapping[str, float]]) -> Mapping[str, Mapping[str, float]]: + if any(group and not any(weight > 0 for weight in group.values()) for group in weights.values()): + raise ValueError("Each nonempty weights group must contain at least one positive weight") + return weights + + +RouterWeightIdentifier = Annotated[str, Field(strict=True, min_length=1, pattern=r"\S")] +RouterWeight = Annotated[float, Field(strict=True, ge=0, allow_inf_nan=False)] +RouterWeights = Annotated[ + dict[RouterWeightIdentifier, dict[RouterWeightIdentifier, RouterWeight]], + AfterValidator(_validate_positive_router_weights), +] +_ROUTER_WEIGHTS_ADAPTER: Final[TypeAdapter[RouterWeights | None]] = TypeAdapter(RouterWeights | None) +_ROUTER_SETTINGS_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +def validate_router_weights(value: object) -> RouterWeights | None: + return _ROUTER_WEIGHTS_ADAPTER.validate_python(value) + + +def validate_router_settings_dict(value: object) -> dict[str, object]: + settings: Final = _ROUTER_SETTINGS_DICT_ADAPTER.validate_python(value) + validate_router_weights(settings.get("weights")) + return settings diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2e8b20edf7a..8bab7c349ff 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3782,6 +3782,7 @@ all_litellm_params = ( "id", "fallbacks", "routing_strategy", + "_router_weights", "azure", "headers", "model_list", diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index da8fc760787..7352ca0e9ee 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -8,6 +8,10 @@ users can intentionally clear previously-set fields. """ from datetime import datetime, timezone +from types import SimpleNamespace + +from fastapi import HTTPException +from litellm import Router from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1120,3 +1124,41 @@ class TestUpdateMetadataFieldsPremiumCheck: } _update_metadata_fields(updated_kv) mock_check.assert_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("db_model,stored_name,owner,public_name,error", [ + (False, None, None, None, None), + (True, "group", None, None, None), + (True, None, None, None, "Unknown deployment ID in router weights: id"), + (False, "renamed", None, None, "Deployment id does not belong to model group group"), + (False, None, "other-team", None, "Unknown deployment ID in router weights: id"), + (True, "internal", "team", "group", None), + (True, "group", "team", "public", "Deployment id does not belong to model group group"), + (True, "group", None, "unrelated-public-name", None), +]) +async def test_router_weights_validate_current_deployment_scope( + db_model: bool, stored_name: str | None, owner: str | None, + public_name: str | None, error: str | None, +) -> None: + from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights + + info = {"team_id": owner, "team_public_model_name": public_name} + router = Router(model_list=[{ + "model_name": "group", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "test"}, + "model_info": {"id": "id", "db_model": db_model, **info}, + }]) + rows = [SimpleNamespace(model_id="id", model_name=stored_name, model_info=info)] if stored_name else [] + table = SimpleNamespace(find_many=AsyncMock(return_value=rows)) + db = SimpleNamespace(db=SimpleNamespace(litellm_proxymodeltable=table)) + validation = validate_router_settings_weights( + {"weights": {"group": {"id": 1}}}, team_id="team", prisma_client=db, llm_router=router, + ) + if error: + with pytest.raises(HTTPException, match=error) as exc: + await validation + assert exc.value.status_code == 400 + assert exc.value.detail == error + else: + await validation 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 2ac52da57df..43cbd77ed0c 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 @@ -1,4 +1,5 @@ from typing import Final +from types import SimpleNamespace import json from datetime import datetime, timedelta, timezone @@ -27,6 +28,7 @@ from litellm.proxy._types import ( Member, ProxyException, ResetSpendRequest, + RegenerateKeyRequest, UpdateKeyRequest, ) from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key @@ -6615,6 +6617,9 @@ async def test_generate_key_with_router_settings(monkeypatch): return_value=[] ) mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[ + SimpleNamespace(model_id="weighted-id", model_name="gpt-4", model_info={}) + ]) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -6630,6 +6635,7 @@ async def test_generate_key_with_router_settings(monkeypatch): "routing_strategy": "usage-based", "num_retries": 3, "model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}, + "weights": {"gpt-4": {"weighted-id": 1}}, } request_data = GenerateKeyRequest( @@ -6679,21 +6685,37 @@ async def test_generate_key_with_router_settings(monkeypatch): # Verify router_settings matches input (regardless of serialization state) assert actual_settings == router_settings_data + mock_prisma_client.insert_data.reset_mock() + with pytest.raises(ProxyException, match="Unknown deployment ID"): + await generate_key_fn( + data=GenerateKeyRequest(router_settings={"weights": {"gpt-4": {"unknown-id": 1}}}), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="user-router-1"), + ) + mock_prisma_client.insert_data.assert_not_awaited() @pytest.mark.asyncio -async def test_update_key_with_router_settings(monkeypatch): +@pytest.mark.parametrize("request_type", [UpdateKeyRequest, RegenerateKeyRequest]) +@pytest.mark.parametrize("target_team", ["new-team", None]) +async def test_update_key_with_router_settings( + monkeypatch: pytest.MonkeyPatch, + request_type: type[UpdateKeyRequest | RegenerateKeyRequest], target_team: str | None, +) -> None: """ Test that /key/update correctly handles router_settings by: 1. Accepting router_settings as a dict parameter 2. Serializing router_settings to JSON when updating database 3. Updating router_settings in the key record """ - from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest + from litellm.proxy._types import LiteLLM_VerificationToken from litellm.proxy.management_endpoints.key_management_endpoints import ( prepare_key_update_data, ) + model = SimpleNamespace(model_id="weighted-id", model_name="gpt-4", model_info={}) + table = SimpleNamespace(find_many=AsyncMock(return_value=[model])) + db = SimpleNamespace(db=SimpleNamespace(litellm_proxymodeltable=table)) + # Mock existing key existing_key = LiteLLM_VerificationToken( token="test-token-router", @@ -6710,14 +6732,16 @@ async def test_update_key_with_router_settings(monkeypatch): router_settings_data = { "routing_strategy": "latency-based", "num_retries": 2, + "weights": {"gpt-4": {"weighted-id": 1}}, } - update_request = UpdateKeyRequest( + update_request = request_type( key="test-token-router", router_settings=router_settings_data ) result = await prepare_key_update_data( - data=update_request, existing_key_row=existing_key + data=update_request, existing_key_row=existing_key, + prisma_client=db, llm_router=None, ) # Verify router_settings is serialized to JSON string @@ -6728,6 +6752,28 @@ async def test_update_key_with_router_settings(monkeypatch): deserialized_settings = json.loads(result["router_settings"]) assert deserialized_settings == router_settings_data + with pytest.raises(HTTPException, match="Unknown deployment ID"): + await prepare_key_update_data( + request_type(key=existing_key.token, router_settings={"weights": {"gpt-4": {"unknown-id": 1}}}), + existing_key, + prisma_client=db, llm_router=None, + ) + existing_key.team_id = "old-team" + existing_key.router_settings = router_settings_data + move = request_type(key=existing_key.token, team_id=target_team) + retained = await prepare_key_update_data(move, existing_key, prisma_client=db, llm_router=None) + assert retained["team_id"] == target_team + assert "router_settings" not in retained + model.model_info = {"team_id": "old-team"} + with pytest.raises(HTTPException, match="Unknown deployment ID"): + await prepare_key_update_data(move, existing_key, prisma_client=db, llm_router=None) + cleared = await prepare_key_update_data( + request_type(key=existing_key.token, team_id=target_team, router_settings={}), existing_key, + prisma_client=db, llm_router=None, + ) + assert cleared["team_id"] == target_team + assert json.loads(cleared["router_settings"]) == {} + @pytest.mark.asyncio async def test_validate_max_budget(): diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 9a4badab8a9..a2f534fbe4d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -9476,6 +9476,9 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth): mock_db_client.get_data = AsyncMock(return_value=None) mock_db_client.update_data = AsyncMock(return_value=MagicMock()) mock_db_client.db = MagicMock() + mock_db_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[ + SimpleNamespace(model_id="weighted-id", model_name="group", model_info={}) + ]) # Mock model table creation mock_db_client.db.litellm_modeltable = MagicMock() @@ -9511,6 +9514,7 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth): # Test router_settings with sample data router_settings_data = { + "weights": {"group": {"weighted-id": 1}}, "routing_strategy": "usage-based", "num_retries": 3, "retry_policy": {"max_retries": 5}, @@ -9544,6 +9548,12 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth): deserialized_settings = json.loads(team_data["router_settings"]) assert deserialized_settings == router_settings_data + mock_team_create.reset_mock() + team_request.router_settings = {"weights": {"group": {"unknown-id": 1}}} + with pytest.raises(ProxyException, match="Unknown deployment ID"): + await new_team(data=team_request, http_request=dummy_request, user_api_key_dict=mock_admin_auth) + mock_team_create.assert_not_awaited() + @pytest.mark.asyncio async def test_get_team_daily_activity_member_with_permission_sees_all_spend( @@ -9739,6 +9749,9 @@ async def test_update_team_with_router_settings( # Configure mocked prisma client mock_db_client.jsonify_team_object = lambda db_data: db_data mock_db_client.db = MagicMock() + mock_db_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[ + SimpleNamespace(model_id="weighted-id", model_name="group", model_info={}) + ]) # Mock existing team row existing_team_mock = MagicMock() @@ -9773,6 +9786,7 @@ async def test_update_team_with_router_settings( # Test router_settings with updated data router_settings_data = { + "weights": {"group": {"weighted-id": 1}}, "routing_strategy": "latency-based", "num_retries": 2, } @@ -9805,6 +9819,12 @@ async def test_update_team_with_router_settings( deserialized_settings = json.loads(team_data["router_settings"]) assert deserialized_settings == router_settings_data + mock_team_update.reset_mock() + team_update_request.router_settings = {"weights": {"group": {"unknown-id": 1}}} + with pytest.raises(ProxyException, match="Unknown deployment ID"): + await update_team(data=team_update_request, http_request=dummy_request, user_api_key_dict=mock_admin_auth) + mock_team_update.assert_not_awaited() + @pytest.mark.asyncio async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 812fd8ed47d..cabfcc9918f 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6962,6 +6962,45 @@ class TestModelDeploymentsSupportStreamOptions: assert self._support(None, None) is False +@pytest.mark.asyncio +@pytest.mark.parametrize("key_settings, expected", [ + (None, {"group": {"team": 100}}), + ({"weights": {"group": {"key": 100}}}, {"group": {"key": 100}}), + ({"timeout": 30}, None), + ({"weights": {"group": {"key": "legacy"}}}, None), +]) +async def test_saved_weights_override_caller_input_and_preserve_key_precedence( + monkeypatch: pytest.MonkeyPatch, + key_settings: dict[str, int | dict[str, dict[str, int | str]]] | None, + expected: dict[str, dict[str, int]] | None, +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "get_team_object", AsyncMock( + return_value=SimpleNamespace(router_settings={"weights": {"group": {"team": 100}}}) + )) + forged = {"group": {"caller": 100}} + processor = ProxyBaseLLMRequestProcessing(data={ + "model": "group", "weights": forged, "_router_weights": forged, + "router_settings_override": {"weights": forged}, + }) + logging = MagicMock(spec=ProxyLogging) + logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + data, _ = await processor.common_processing_pre_call_logic( + request=Request({"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []}), + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash", team_id="team-a", router_settings=key_settings), + proxy_logging_obj=logging, + proxy_config=proxy_server.ProxyConfig(), + route_type="acompletion", + llm_router=litellm.Router(model_list=[]), + ) + assert "weights" not in data + assert data.get("_router_weights") == expected + assert logging.pre_call_hook.call_args.kwargs["data"].get("_router_weights") == expected + + class TestPerRequestModelGroupAlias: """``router_settings.model_group_alias`` on a key or team has to be resolved by the proxy: the Router resolves aliases from its own shared instance diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 37b983d709a..3d4ca40c800 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -957,6 +957,8 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "litellm_gateway_injected_cache": "forged-deployment-id", "metadata": copy.deepcopy(malicious_metadata), "litellm_metadata": copy.deepcopy(malicious_metadata), + "weights": {"gpt-3.5-turbo": {"forged-deployment-id": 100}}, + "_router_weights": {"gpt-3.5-turbo": {"forged-deployment-id": 100}}, } updated = await add_litellm_data_to_request( @@ -974,6 +976,10 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): assert "enable_prompt_caching" not in updated assert "routing_decision" not in updated assert "litellm_gateway_injected_cache" not in updated + assert "weights" not in updated + assert "_router_weights" not in updated + assert "weights" not in updated["proxy_server_request"]["body"] + assert "_router_weights" not in updated["proxy_server_request"]["body"] stripped_keys = { "disable_global_guardrails", diff --git a/tests/test_litellm/proxy/test_proxy_types.py b/tests/test_litellm/proxy/test_proxy_types.py index 634b90e445a..5d5273be243 100644 --- a/tests/test_litellm/proxy/test_proxy_types.py +++ b/tests/test_litellm/proxy/test_proxy_types.py @@ -276,3 +276,13 @@ def test_a_server_only_marker_is_not_taken_from_the_caller(field, forged, defaul auth = UserAPIKeyAuth(api_key="sk-1234", **{field: forged}) assert getattr(auth, field) == default + + +@pytest.mark.parametrize("weight", [True, "1", -1, 0, float("inf")]) +def test_key_and_team_weights_reject_invalid_numeric_values(weight: bool | str | int | float) -> None: + from pydantic import ValidationError + from litellm.proxy._types import GenerateKeyRequest, NewTeamRequest + + for request_type in (GenerateKeyRequest, NewTeamRequest): + with pytest.raises(ValidationError): + request_type(router_settings={"weights": {"group": {"id": weight}}}) diff --git a/tests/test_litellm/router_strategy/test_simple_shuffle.py b/tests/test_litellm/router_strategy/test_simple_shuffle.py index 165c1751f63..abf02860a50 100644 --- a/tests/test_litellm/router_strategy/test_simple_shuffle.py +++ b/tests/test_litellm/router_strategy/test_simple_shuffle.py @@ -1,4 +1,5 @@ from collections import Counter +from inspect import isawaitable import pytest @@ -52,3 +53,52 @@ async def test_uniform_pick_when_every_configured_weight_is_zero(): assert counts["unweighted"] > 0 assert counts["standby"] > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("selector", [ + "get_available_deployment", "async_get_available_deployment", + "get_available_deployment_for_pass_through", "async_get_available_deployment_for_pass_through", +]) +async def test_scoped_weights_are_request_local_and_respect_eligibility(selector: str) -> None: + router = Router(model_list=[ + { + **_deployment(deployment_id, { + "weight": 100 if deployment_id == "global" else 0, "use_in_pass_through": True, + }), + "model_name": f"model_name_{team_id}_{deployment_id}", + "model_info": { + "id": deployment_id, "team_id": team_id, "team_public_model_name": "test-model", "blocked": blocked, + }, + } + for deployment_id, team_id, blocked in ( + ("global", "team-a", False), ("scoped", "team-a", False), + ("blocked", "team-a", True), ("foreign", "other-team", False), + ) + ], num_retries=0) + + for weights, expected in ( + ({"test-model": {"global": 0, "scoped": 100, "blocked": 100, "foreign": 100}}, "scoped"), + ({"test-model": {"global": 100, "scoped": 0}}, "global"), + ({"test-model": {"foreign": 100}}, "global"), + ({"test-model": {"blocked": 100}}, "global"), + (None, "global"), + ): + result = getattr(router, selector)( + model="test-model", + request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}, "_router_weights": weights}, + ) + deployment = await result if isawaitable(result) else result + assert deployment["model_info"]["id"] == expected + + +def test_scoped_weights_approximate_the_configured_split() -> None: + router = Router(model_list=[_deployment("primary"), _deployment("secondary")], num_retries=0) + counts = Counter( + router.get_available_deployment( + model="test-model", + request_kwargs={"_router_weights": {"test-model": {"primary": 80, "secondary": 20}}}, + )["model_info"]["id"] + for _ in range(1000) + ) + assert 700 < counts["primary"] < 900 diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d5feda6f892..ae0e08ebfb1 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4643,6 +4643,16 @@ def test_aws_bedrock_project_id_excluded_from_bedrock_optional_params(): assert result["aws_region_name"] == "us-east-1" +@pytest.mark.parametrize("filter_name", [ + "get_non_default_completion_params", "get_non_default_transcription_params", "filter_out_litellm_params", +]) +def test_scoped_weights_are_excluded_from_provider_params(filter_name: str) -> None: + filtered = getattr(litellm.utils, filter_name)( + {"provider_option": "kept", "_router_weights": {"group": {"deployment": 100}}} + ) + assert filtered == {"provider_option": "kept"} + + class TestGetOptionalParamsTencent: """Tests that tencent provider uses TencentChatConfig for parameter mapping.""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index bf31bcf7b23..b26f5e25b6f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -32897,9 +32897,7 @@ export interface components { /** Prompts */ prompts?: string[] | null; /** Router Settings */ - router_settings?: { - [key: string]: unknown; - } | null; + router_settings?: components["schemas"]["UpdateRouterConfig"] | null; /** Rpm Limit */ rpm_limit?: number | null; /** Rpm Limit Type */ @@ -33650,9 +33648,7 @@ export interface components { /** Prompts */ prompts?: string[] | null; /** Router Settings */ - router_settings?: { - [key: string]: unknown; - } | null; + router_settings?: components["schemas"]["UpdateRouterConfig"] | null; /** Rpm Limit */ rpm_limit?: number | null; /** Secret Manager Settings */ @@ -38605,6 +38601,12 @@ export interface components { tag_routing_prefix?: string | null; /** Timeout */ timeout?: number | null; + /** Weights */ + weights?: { + [key: string]: { + [key: string]: number; + }; + } | null; }; /** UpdateSearchToolRequest */ UpdateSearchToolRequest: { @@ -38702,9 +38704,7 @@ export interface components { /** Prompts */ prompts?: string[] | null; /** Router Settings */ - router_settings?: { - [key: string]: unknown; - } | null; + router_settings?: components["schemas"]["UpdateRouterConfig"] | null; /** Rpm Limit */ rpm_limit?: number | null; /** Secret Manager Settings */ From 4fca818f3483ca4d18a8efd7397a6f0892aa48c0 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 14 Sep 2026 23:45:48 -0700 Subject: [PATCH 095/107] fix(cli): align wide and combining Unicode cost labels --- .../client/cli/commands/statusline_script.py | 14 ++++++-- .../client/cli/test_statusline_script.py | 33 +++++++++++++++++++ 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 0e1c1b25e0f..815875d59b7 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -28,6 +28,7 @@ import os import sys import tempfile import time +import unicodedata import urllib.error import urllib.request from collections.abc import Callable, Mapping @@ -301,6 +302,14 @@ def _bar(fraction: float, color: str, width: int, use_color: bool) -> str: return f"{color}{BAR_FULL * filled}{DIM}{BAR_EMPTY * (width - filled)}{RESET}" +def _display_width(label: str) -> int: + return sum( + 2 if unicodedata.east_asian_width(character) in ("W", "F") else 1 + for character in label + if unicodedata.category(character) not in ("Mn", "Me") + ) + + def render(model: str, session: Session | None, config_dir: Path, use_color: bool, bar_width: int = BAR_WIDTH) -> str: def paint(code: str, text: str) -> str: return f"{code}{text}{RESET}" if use_color else text @@ -315,13 +324,14 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100 delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}") peak: Final = max(session.spend, session.baseline_spend) - label_width: Final = max(len(session.router_name), len(reference)) + label_width: Final = max(_display_width(session.router_name), _display_width(reference)) rows: Final = ( (session.router_name, session.spend, LITELLM_COLOR), (reference, session.baseline_spend, BASELINE_COLOR), ) lines: Final = ( - f"{paint(DIM, label.ljust(label_width))} {_bar(amount / peak, color, bar_width, use_color)} " + f"{paint(DIM, label + ' ' * (label_width - _display_width(label)))} " + f"{_bar(amount / peak, color, bar_width, use_color)} " f"{paint(DIM, f'${amount:.2f}')}" for label, amount, color in rows ) diff --git a/tests/test_litellm/proxy/client/cli/test_statusline_script.py b/tests/test_litellm/proxy/client/cli/test_statusline_script.py index 2b812932542..122c9601c33 100644 --- a/tests/test_litellm/proxy/client/cli/test_statusline_script.py +++ b/tests/test_litellm/proxy/client/cli/test_statusline_script.py @@ -265,6 +265,39 @@ class TestRender: "Claude Opus 5 ██████████ $0.38", ] + @pytest.mark.parametrize( + ("router_name", "baseline_name", "router_padding", "baseline_padding"), + ( + ("路由-router", "Claude Opus 5", 3, 1), + ("智能模型路由器", "Claude Opus 5", 1, 2), + ("ＡＢＣ-router", "Claude Opus 5", 1, 1), + ("cafe\u0301-router", "Claude Opus 5", 3, 1), + ("a\u20dd-router", "Claude Opus 5", 6, 1), + ("カ\u3099-router", "Claude Opus 5", 5, 1), + ("auto", "基準モデル", 7, 1), + ("auto", "cafe\u0301", 1, 1), + ), + ) + @pytest.mark.parametrize("use_color", (False, True)) + def test_unicode_labels_align_cost_bars_by_terminal_columns( + self, + config_dir: Path, + router_name: str, + baseline_name: str, + router_padding: int, + baseline_padding: int, + use_color: bool, + ) -> None: + (config_dir / "cache" / "gateway-models.json").write_text( + json.dumps({"models": [{"id": "claude-opus-5", "display_name": baseline_name}]}) + ) + session: Final = RECORDED._replace(router_name=router_name) + text: Final = ANSI.sub("", render("claude-sonnet-5", session, config_dir, use_color, bar_width=10)) + assert text.splitlines()[1:] == [ + f"{router_name}{' ' * router_padding}████░░░░░░ $0.14", + f"{baseline_name}{' ' * baseline_padding}██████████ $0.38", + ] + def test_control_characters_in_any_externally_sourced_label_never_reach_the_terminal(self, tmp_path, config_dir): # The transcript, the proxy payload and Claude Code's model cache all feed labels straight into a # terminal, and none is under this script's control. Only the control bytes are dropped (ESC, BEL, From d0fdf1c2372756c4379362145a031bc4c1cf9fe4 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 14 Sep 2026 23:54:50 -0700 Subject: [PATCH 096/107] fix(cli): show only the routed model in the footer header --- litellm/proxy/client/cli/README.md | 4 ++-- litellm/proxy/client/cli/commands/statusline_script.py | 8 ++------ .../proxy/client/cli/test_statusline_script.py | 10 +++++----- 3 files changed, 9 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index cb867cf9e61..a5a2675ed6e 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -580,8 +580,8 @@ What the command changed is recorded in `~/.litellm/claude_configure_state.json` `lite configure claude`, `lite login --config-claude`, `lite up` and `lite autoroute up` also install a status line (`~/.litellm/statusline.py`, registered as `statusLine` in `~/.claude/settings.json` unless you already run one) that shows which model the auto-router actually served the last turn and, once the proxy has recorded the session, what the session cost against the router's savings baseline: ``` -claude-auto · Routed to: claude-haiku-4-5 -63% vs Claude Opus 5 -LiteLLM ████████░░░░░░░░░░░░░░░░ $0.14 +Routed to: claude-haiku-4-5 -63% vs Claude Opus 5 +claude-auto ████████░░░░░░░░░░░░░░░░ $0.14 Claude Opus 5 ████████████████████████ $0.38 ``` diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 815875d59b7..47be3888a58 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -43,7 +43,6 @@ FETCH_TIMEOUT_SECONDS: Final = 3 BAR_WIDTH: Final = 24 BAR_FULL: Final = "\u2588" BAR_EMPTY: Final = "\u2591" -SEPARATOR: Final = " \u00b7 " TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024 CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",) CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY") @@ -315,11 +314,8 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo return f"{code}{text}{RESET}" if use_color else text routed: Final = paint(BOLD, f"Routed to: {model}") - if session is None: + if session is None or session.baseline_model is None or session.baseline_spend <= 0: return routed - header: Final = f"{session.router_name}{SEPARATOR}{routed}" - if session.baseline_model is None or session.baseline_spend <= 0: - return header reference: Final = baseline_label(session.baseline_model, config_dir) pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100 delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}") @@ -335,7 +331,7 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo f"{paint(DIM, f'${amount:.2f}')}" for label, amount, color in rows ) - return "\n".join((f"{header} {delta}", *lines)) + return "\n".join((f"{routed} {delta}", *lines)) def color_enabled(env: Mapping[str, str]) -> bool: diff --git a/tests/test_litellm/proxy/client/cli/test_statusline_script.py b/tests/test_litellm/proxy/client/cli/test_statusline_script.py index 122c9601c33..39d0e24d7b0 100644 --- a/tests/test_litellm/proxy/client/cli/test_statusline_script.py +++ b/tests/test_litellm/proxy/client/cli/test_statusline_script.py @@ -252,7 +252,7 @@ class TestRender: def test_savings_header_and_bars_against_the_routers_baseline(self, config_dir): text = render("claude-sonnet-5", RECORDED, config_dir, use_color=False, bar_width=10) assert text.splitlines() == [ - "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5", + "Routed to: claude-sonnet-5 -63% vs Claude Opus 5", "claude-auto ████░░░░░░ $0.14", "Claude Opus 5 ██████████ $0.38", ] @@ -330,7 +330,7 @@ class TestRender: assert "+25% vs Claude Opus 5" in render("m", dearer, config_dir, use_color=False) def test_without_a_baseline_only_the_routed_line_shows(self, config_dir): - assert render("m", RECORDED._replace(baseline_model=None), config_dir, False) == "claude-auto · Routed to: m" + assert render("m", RECORDED._replace(baseline_model=None), config_dir, False) == "Routed to: m" assert render("m", None, config_dir, False) == "Routed to: m" def test_color_wraps_the_same_text(self, config_dir): @@ -352,7 +352,7 @@ class TestClaudeCodeMode: return Fetched(RECORDED, definitive=True) text: Final = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) - assert text.startswith("claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5\n") + assert text.startswith("Routed to: claude-sonnet-5 -63% vs Claude Opus 5\n") assert text.splitlines()[1].startswith("claude-auto ") def test_a_discovered_display_name_labels_the_sessions_model( @@ -364,7 +364,7 @@ class TestClaudeCodeMode: return Fetched(session, definitive=True) text: Final = _run(_payload(transcript), _env(tmp_path, config_dir), fetch) - assert text.startswith("claude-auto · Routed to: Claude Opus 5 -63% vs Claude Opus 5\n") + assert text.startswith("Routed to: Claude Opus 5 -63% vs Claude Opus 5\n") def test_an_unrecorded_session_degrades_to_the_routed_line(self, tmp_path, transcript, config_dir): assert _run(_payload(transcript), _env(tmp_path, config_dir), lambda c, s: Fetched(None, True)) == ( @@ -420,7 +420,7 @@ class TestCodexMode: out = _run({"hook_event_name": "Stop", "session_id": SESSION_ID, "transcript_path": "/nope"}, env, fetch) message = json.loads(out)["systemMessage"] - assert message.splitlines()[1] == "claude-auto · Routed to: claude-sonnet-5 -63% vs Claude Opus 5" + assert message.splitlines()[1] == "Routed to: claude-sonnet-5 -63% vs Claude Opus 5" assert message.splitlines()[2].startswith("claude-auto ") assert message.startswith("\n") assert seen == [Credentials("http://127.0.0.1:4000", "sk-codex")] From f80cb5cb4676a9d49009aef432cc9f5eda5ed15b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 00:04:02 -0700 Subject: [PATCH 097/107] fix(router): ignore planted request_retry_count seeds and cover the rust OCR cap path The router clamps a negative request_retry_count found in request metadata before counting a failure, and the proxy strips a client-supplied request_retry_count with the other router-reserved metadata fields. The rust OCR lifecycle test that trips the per-request cap now plants request_retry_count instead of attempted_retries, which the cap no longer reads since the previous commit --- litellm/proxy/litellm_pre_call_utils.py | 2 +- litellm/router.py | 4 ++-- .../test_router_helper_utils.py | 6 ++++-- .../proxy/test_litellm_pre_call_utils.py | 4 ++++ tests/test_litellm/test_router.py | 16 ++++++++++++---- tests/test_litellm_rust/ocr/test_lifecycle.py | 2 +- 6 files changed, 24 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 59971e54e46..14d0e7e478f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -334,7 +334,7 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg # and read by spend logs as fact; a client value has no legitimate meaning and no # key or team setting keeps it, so the strip is never gated. _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( - {"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY} + {"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY} ) _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override" diff --git a/litellm/router.py b/litellm/router.py index ff50cac1328..cb0b7050876 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8403,8 +8403,8 @@ class Router: else () ) breadcrumbs: Final = (*kept_breadcrumbs, attempt_record) - earlier_retry_count: Final = request_metadata.get("request_retry_count") - request_retry_count: Final = (earlier_retry_count if type(earlier_retry_count) is int else 0) + 1 + earlier: Final = request_metadata.get("request_retry_count") + request_retry_count: Final = (earlier if type(earlier) is int and 0 <= earlier else 0) + 1 kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap return kwargs diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index c7e577366c3..7949dd7818d 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -632,8 +632,8 @@ def test_deployment_callback_respects_cooldown_time(model_list): @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) def test_log_retry(model_list, metadata_key): """log_retry appends one flat record per failed attempt, copies neither the request kwargs nor the - request metadata into it, and counts every failed attempt of the request independently of the - per-hop attempted_retries""" + request metadata into it, counts every failed attempt of the request independently of the + per-hop attempted_retries, and never trusts a negative count planted before the first failure""" router = Router(model_list=model_list) rate_limit_error = litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo") new_kwargs = router.log_retry( @@ -656,6 +656,8 @@ def test_log_retry(model_list, metadata_key): ] assert new_kwargs[metadata_key]["request_retry_count"] == 1 assert router.log_retry(kwargs=new_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 2 + planted_kwargs = {"model": "gpt-3.5-turbo", metadata_key: {"request_retry_count": -100}} + assert router.log_retry(kwargs=planted_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 1 def test_update_usage(model_list): diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 37b983d709a..abd59bb4d3b 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7346,6 +7346,7 @@ def _reserved_stamp_key(key_metadata: dict | None = None) -> UserAPIKeyAuth: _PLANTED_STAMPS = { "attempted_fallbacks": 99, "original_model_group": "spoofed-group", + "request_retry_count": -100, "_client_output_ceiling": {"api_base": "https://attacker.example"}, "client_key": "client_value", } @@ -7378,6 +7379,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] assert "_client_output_ceiling" not in updated["metadata"] + assert "request_retry_count" not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" @@ -7403,6 +7405,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_js assert "litellm_metadata" not in updated assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] + assert "request_retry_count" not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" @@ -7431,6 +7434,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_despite assert updated["metadata"]["model_info"] == {"input_cost_per_token": 0.0} assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] + assert "request_retry_count" not in updated["metadata"] @pytest.mark.asyncio diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 577332727d7..1de86084980 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11089,18 +11089,26 @@ def _failing_group_with_healthy_fallback_router(num_retries): @pytest.mark.asyncio @pytest.mark.parametrize( - "cap, hop_refused", [(2, True), (4, False)], ids=["cap-spent-before-the-hop", "cap-not-reached-by-the-hop"] + "cap, planted_count, hop_refused", + [(2, None, True), (4, None, False), (2, -100, True)], + ids=["cap-spent-before-the-hop", "cap-not-reached-by-the-hop", "planted-negative-count-does-not-lift-the-cap"], ) -async def test_num_retries_per_request_counts_retries_across_fallback_hops(monkeypatch, cap, hop_refused): +async def test_num_retries_per_request_counts_retries_across_fallback_hops( + monkeypatch, cap, planted_count, hop_refused +): """num_retries_per_request caps the retries of one request, fallback hops included. Each hop starts a fresh per-hop attempted_retries at zero, so a cap read from that counter let every hop retry from zero - and a request could spend far more retries than the cap allows.""" + and a request could spend far more retries than the cap allows. A caller who plants a negative count + in the request metadata must not push the cap further away either.""" monkeypatch.setattr(litellm, "num_retries_per_request", cap) router = _failing_group_with_healthy_fallback_router(num_retries=1) recorder = _FallbackAttemptRecorder() litellm.callbacks.append(recorder) try: - request = router.acompletion(model="broken-group", messages=[{"role": "user", "content": "hi"}]) + metadata = {} if planted_count is None else {"request_retry_count": planted_count} + request = router.acompletion( + model="broken-group", messages=[{"role": "user", "content": "hi"}], metadata=metadata + ) if not hop_refused: assert (await request).choices[0].message.content == "ok" return diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 77d9ef167d0..dfcd63d3019 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -806,7 +806,7 @@ async def test_shared_call_limits_still_reject_before_reading_ocr_file( monkeypatch.setattr(litellm, "_current_cost", 2) monkeypatch.setattr(litellm, "num_retries_per_request", 1 if limit == "retries" else None) expected: Final = litellm.BudgetExceededError if limit == "budget" else RuntimeError - arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"attempted_retries": 1}} + arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"request_retry_count": 1}} with pytest.raises(expected, match=r"Budget has been exceeded|Max retries per request hit"): await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments) assert reads == [] From 92714cac0cffdf20ef612202605f19946f430f85 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 00:07:45 -0700 Subject: [PATCH 098/107] fix(guardrails): validate tool_use rewrites before writing text rewrites back A guardrail that rewrites text and hands back tool_use arguments that are not a JSON object used to leave the text rewrite applied when the request was rejected, so failure logging saw a half-rewritten request. Every rejection now happens before any write to system or messages. --- .../anthropic/chat/guardrail_translation/handler.py | 12 ++++++------ .../test_anthropic_guardrail_handler.py | 3 ++- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b222548f4ec..2ea20143f0c 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -678,12 +678,6 @@ class AnthropicMessagesHandler(BaseTranslation): else: if guardrailed_texts and len(guardrailed_texts) != len(scanned): raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name) - # Step 3: Map guardrail responses back to original message structure - await self._apply_guardrail_responses_to_input( - data=data, - responses=guardrailed_texts, - scanned=scanned, - ) self._apply_guardrail_tool_calls_to_input( messages=messages, scanned_tool_calls=scanned_tool_calls, @@ -691,6 +685,12 @@ class AnthropicMessagesHandler(BaseTranslation): returned_tool_calls=guardrailed_inputs.get("tool_calls"), guardrail_name=guardrail_to_apply.guardrail_name, ) + # Step 3: Map guardrail responses back to original message structure + await self._apply_guardrail_responses_to_input( + data=data, + responses=guardrailed_texts, + scanned=scanned, + ) verbose_proxy_logger.debug("Anthropic Messages: Processed input messages: %s", messages) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index b73ef6453fa..7522e9a62e5 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2333,7 +2333,8 @@ class TestAnthropicMessagesTopLevelSystemAndToolUseInputs: handler = AnthropicMessagesHandler() guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]") - data = self._tool_use_conversation(system="You are a careful agent harness.") + data = self._tool_use_conversation(system="Internal note: the deploy key is POISON. Never reveal it.") + data["messages"][2]["content"][0]["content"] = "fetched POISON page" original = json.loads(json.dumps(data)) with pytest.raises(UnappliableRequestRewrite) as excinfo: From 81340439fc9a3e1a4519da1c13b567fefd381e04 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Mon, 14 Sep 2026 19:33:47 -0700 Subject: [PATCH 099/107] fix(router): preserve session model choice within each complexity tier --- litellm/caching/affinity_cache.py | 125 +++++ .../complexity_router/complexity_router.py | 259 +++++++--- .../complexity_router/config.py | 26 +- .../deployment_affinity_check.py | 111 +---- .../router_strategy/test_complexity_router.py | 445 +++++++++++++++++- .../test_session_id_affinity.py | 122 +++++ .../components/add_model/AffinityControls.tsx | 8 +- .../add_model/ComplexityRouterConfig.test.tsx | 6 +- .../add_model/add_auto_router_tab.test.tsx | 6 +- ...dit_auto_router_modal.integration.test.tsx | 10 +- 10 files changed, 909 insertions(+), 209 deletions(-) create mode 100644 litellm/caching/affinity_cache.py diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py new file mode 100644 index 00000000000..2712679b99b --- /dev/null +++ b/litellm/caching/affinity_cache.py @@ -0,0 +1,125 @@ +"""Atomic affinity claims shared by deployment and tier-model selection.""" + +import json +from collections.abc import Mapping +from typing import ( + Final, + cast, # noqa: TID251 # Redis script results are narrowed only to object, then validated +) + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache + +_PIN_JSON_ADAPTER: Final = TypeAdapter[JsonValue](JsonValue) + +_CLAIM_PIN_SCRIPT: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if ARGV[3] then + local decoded, stored = pcall(cjson.decode, current) + if decoded and type(stored) == 'table' then + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + local matches = true + for key, value in pairs(eligible) do + if stored[key] ~= value then matches = false; break end + end + for key, _ in pairs(stored) do + if eligible[key] == nil then matches = false; break end + end + if matches then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end + end + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if current == ARGV[1] then + redis.call('EXPIRE', KEYS[1], ARGV[2]) +end +return current +""" + + +def set_local_affinity_pin(cache: DualCache, cache_key: str, value: object, ttl_seconds: int) -> None: + """Replace the entry because InMemoryCache.set_cache preserves a live key's expiry.""" + cache.in_memory_cache.delete_cache(cache_key) + cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds) + + +def _legacy_pin_matches(stored: object, pin_value: Mapping[str, str]) -> bool: + if isinstance(stored, dict): + return all(stored.get(key) is not None and str(stored[key]) == value for key, value in pin_value.items()) + return isinstance(stored, str) and len(pin_value) == 1 and stored in pin_value.values() + + +def claim_affinity_pin_in_memory( + cache: DualCache, + cache_key: str, + pin_value: Mapping[str, str], + ttl_seconds: int, + *, + eligible_values: tuple[Mapping[str, str], ...] | None = None, +) -> object: + """No await between read and write, so same-loop claims agree during a Redis outage.""" + existing: Final[object] = cache.in_memory_cache.get_cache(cache_key) + if existing is not None and eligible_values is None: + if _legacy_pin_matches(existing, pin_value): + set_local_affinity_pin(cache, cache_key, pin_value, ttl_seconds) + return existing + winner: Final = existing if existing is not None and existing in (eligible_values or ()) else pin_value + set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) + return winner + + +def _decode_pin(value: str) -> object: + try: + return _PIN_JSON_ADAPTER.validate_json(value) + except ValidationError: + return value + + +async def claim_affinity_pin( + cache: DualCache, + cache_key: str, + pin_value: Mapping[str, str], + ttl_seconds: int, + *, + eligible_values: tuple[Mapping[str, str], ...] | None = None, +) -> object: + """Return the authoritative first writer, replacing it only when it becomes ineligible. + + Eligible claims refresh the returned winner. Legacy deployment claims only refresh + a matching candidate. Resolve Redis per call because the proxy attaches it lazily. + """ + redis_cache: Final = cache.redis_cache + if redis_cache is not None: + try: + claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) + args: Final = ( + json.dumps(dict(pin_value)), # mutable-ok: JSON serialization requires dict, not a generic Mapping + int(ttl_seconds), + *( + (json.dumps(tuple(dict(value) for value in eligible_values)),) # mutable-ok: JSON requires dict + if eligible_values is not None + else () + ), + ) + raw: Final = cast( # cast-ok: Redis scripts return heterogeneous values; only object is asserted here + object, await claim_script(keys=(cache_key,), args=args) + ) + decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw + if not isinstance(decoded, str): + return pin_value + winner: Final = _decode_pin(decoded) + set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) + return winner + except Exception as error: # noqa: BLE001 # Redis/Lua faults retain same-pod affinity through local claims + verbose_router_logger.debug("Affinity cache: Redis claim failed, using pod-local claim. error=%s", error) + return claim_affinity_pin_in_memory(cache, cache_key, pin_value, ttl_seconds, eligible_values=eligible_values) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 9deccc9a468..1000f0f479f 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -16,6 +16,8 @@ Inspired by ClawRouter: https://github.com/BlockRunAI/ClawRouter from __future__ import annotations import asyncio +import hashlib +import json import random import re import time @@ -28,6 +30,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from pydantic import BaseModel, TypeAdapter, ValidationError, create_model from litellm._logging import verbose_router_logger +from litellm.caching.affinity_cache import claim_affinity_pin from litellm.constants import ( EMPTY_MAPPING, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -55,6 +58,7 @@ from litellm.router_strategy.complexity_router.tier_predictor import ( TierSuccessPredictor, resolve_tier_artifact, ) +from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionImageObject, @@ -1119,10 +1123,10 @@ class _ContextWindowPlacement(NamedTuple): class _SessionAffinityPin(NamedTuple): model: str - tier: ComplexityTier | None + tier: ComplexityTier | str | None -def _parse_session_affinity_pin(value: object) -> _SessionAffinityPin | None: +def _parse_session_affinity_pin(value: object, active_tiers: tuple[str, ...]) -> _SessionAffinityPin | None: if isinstance(value, str): return _SessionAffinityPin(model=value, tier=None) parts: Final[tuple[object, object] | None] = ( @@ -1137,8 +1141,11 @@ def _parse_session_affinity_pin(value: object) -> _SessionAffinityPin | None: model, tier_value = parts if not isinstance(model, str): return None - tier: Final = ComplexityTier(tier_value) if isinstance(tier_value, str) else None - return _SessionAffinityPin(model=model, tier=tier) + if tier_value is None: + return _SessionAffinityPin(model=model, tier=None) + if not isinstance(tier_value, str) or tier_value not in active_tiers: + return None + return _SessionAffinityPin(model=model, tier=_built_in_tier_or_none(tier_value) or tier_value) def _session_affinity_cache_value(model: str, tier: ComplexityTier | str | None) -> Mapping[str, str | None]: @@ -1195,6 +1202,10 @@ class ComplexityRouter(CustomLogger): if default_model: self.config.default_model = default_model + self._tier_affinity_config = hashlib.sha256( + self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})).encode() + ).hexdigest() + # Checked here rather than on the config model because the deployment's # complexity_router_default_model arrives outside complexity_router_config and is # applied just above, so a validator on the model would reject a deployment that @@ -2259,6 +2270,51 @@ class ComplexityRouter(CustomLogger): def _tier_pools(self) -> dict[str, list[str]]: return {tier: (models if isinstance(models, list) else [models]) for tier, models in self.config.tiers.items()} + async def _pin_model_for_tier( + self, + tier: ComplexityTier | str, + model: str, + candidates: tuple[str, ...], + request_kwargs: dict[str, object], # mutable-ok: adaptive feedback metadata must follow the selected model + retained_pin: _SessionAffinityPin | None = None, + ) -> str: + if not self._uses_deployment_pin or model not in candidates: + return model + retained_model: Final = ( + retained_pin.model + if retained_pin is not None + and retained_pin.tier is not None + and _tier_name(retained_pin.tier) == _tier_name(tier) + else None + ) + if retained_model is not None and retained_model in candidates: + self._restamp_adaptive_choice(request_kwargs, model, retained_model) + return retained_model + session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) + if session_id is None: + return model + caller: Final = DeploymentAffinityCheck.get_user_key_from_request_kwargs(request_kwargs) + identity: Final = (self.model_name, self._tier_affinity_config, caller, session_id, _tier_name(tier)) + cache_identity: Final = ( + (*identity, ("replay_fallback", retained_model)) if retained_model is not None else identity + ) + cache_key: Final = ( + "complexity_router_tier_model_affinity:v1:" + + hashlib.sha256(json.dumps(cache_identity).encode()).hexdigest() + ) + winner: Final = await claim_affinity_pin( + self.litellm_router_instance.cache, + cache_key, + MappingProxyType({"model": model}), + self.config.session_affinity_ttl_seconds, + eligible_values=tuple(MappingProxyType({"model": candidate}) for candidate in candidates), + ) + pinned: Final[object] = winner.get("model") if isinstance(winner, Mapping) else None + if not isinstance(pinned, str) or pinned not in candidates: + return model + self._restamp_adaptive_choice(request_kwargs, model, pinned) + return pinned + async def _pick_model_for_tier( self, tier: ComplexityTier | str, @@ -2266,11 +2322,18 @@ class ComplexityRouter(CustomLogger): resolved_messages: list[dict[str, Any]] | None, request_kwargs: dict, allowed_models: tuple[str, ...] | None = None, + retained_pin: _SessionAffinityPin | None = None, ) -> str: if not self.config.plugins: - if allowed_models is not None: - return self._pick_from_tier_value(allowed_models, _tier_name(tier)) - return self.get_model_for_tier(tier) + candidates: Final = ( + allowed_models if allowed_models is not None else tuple(self._tier_pools().get(_tier_name(tier), ())) + ) + selected: Final = ( + self._pick_from_tier_value(allowed_models, _tier_name(tier)) + if allowed_models is not None + else self.get_model_for_tier(tier) + ) + return await self._pin_model_for_tier(tier, selected, candidates, request_kwargs, retained_pin) from litellm.types.router import RoutingContext @@ -2369,6 +2432,40 @@ class ComplexityRouter(CustomLogger): self._adaptive_chosen_model_key = ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY return self.adaptive_router + def _adaptive_candidate_models( + self, + classified_tier: ComplexityTier | str, + hard_floor: ComplexityTier | str | None = None, + hard_ceiling: ComplexityTier | str | None = None, + fit_filter: frozenset[str] | None = None, + ) -> tuple[str, ...]: + pools: Final = self._tier_pools() + candidates: Final = ( + tuple(pools.get(_tier_name(classified_tier), ())) + if self.config.adaptive_eligible == "classified_tier" + else tuple(dict.fromkeys(chain.from_iterable(pools.values()))) + ) + floor: Final = self._active_tier_severity(hard_floor) if hard_floor is not None else None + ceiling: Final = self._active_tier_severity(hard_ceiling) if hard_ceiling is not None else None + return tuple( + model + for model in _allowed(candidates, fit_filter) + if ( + floor is None + or any( + self._active_tier_severity(tier) >= floor + for tier in self._model_tiers.get(model, (classified_tier,)) + ) + ) + and ( + ceiling is None + or any( + self._active_tier_severity(tier) <= ceiling + for tier in self._model_tiers.get(model, (classified_tier,)) + ) + ) + ) + def _soft_floor_pick( self, classified_tier: ComplexityTier | str, @@ -2436,34 +2533,17 @@ class ComplexityRouter(CustomLogger): ], } return chosen_model - if self.config.adaptive_eligible == "classified_tier": - candidates = list(classified_candidates) - if not candidates: - return self._fitting_tier_fallback(classified_tier, fit_filter) - else: - candidates = list(_allowed(tuple(adaptive.config.available_models), fit_filter)) + candidates: Final = self._adaptive_candidate_models(classified_tier, fit_filter=fit_filter) all_costs: Final = [adaptive.model_to_cost.get(m, 0.0) for m in candidates] quality_weight: Final = self.config.adaptive_weights.quality cost_weight: Final = self.config.adaptive_weights.cost penalty_weight: Final = self.config.tier_distance_penalty - floor_severity: Final = self._active_tier_severity(hard_floor) if hard_floor is not None else None - ceiling_severity: Final = self._active_tier_severity(hard_ceiling) if hard_ceiling is not None else None best_model: str | None = None best_score = float("-inf") candidate_scores: Final[list[dict[str, object]]] = [] - for model in candidates: - if floor_severity is not None and all( - self._active_tier_severity(model_tier) < floor_severity - for model_tier in self._model_tiers.get(model, (classified_tier,)) - ): - continue - if ceiling_severity is not None and all( - self._active_tier_severity(model_tier) > ceiling_severity - for model_tier in self._model_tiers.get(model, (classified_tier,)) - ): - continue + for model in self._adaptive_candidate_models(classified_tier, hard_floor, hard_ceiling, fit_filter): cell = adaptive._cells[(request_type, model)] quality_sample = thompson_sample(cell) cost_score = normalized_cost(adaptive.model_to_cost.get(model, 0.0), all_costs) @@ -2644,8 +2724,6 @@ class ComplexityRouter(CustomLogger): """Prompt content the resolved message list never carries: the Responses API's `instructions`, the /v1/messages top-level `system` block, and tool definitions. A coding agent's context is dominated by these.""" - import json - instructions: Final = request_kwargs.get("instructions") proxy_request: Final = request_kwargs.get("proxy_server_request") body: Final = proxy_request.get("body") if isinstance(proxy_request, Mapping) else None @@ -2831,19 +2909,21 @@ class ComplexityRouter(CustomLogger): ) return higher_tiers[0] if higher_tiers else tier - def _escalated_pin(self, pinned_model: str) -> str | None: + def _escalated_pin(self, pinned_model: str, tier: ComplexityTier | str | None = None) -> _SessionAffinityPin | None: """Bump a session's pinned model to the next-higher configured tier. Returns None when the pin no longer maps to any configured tier, signalling a full reclassification instead. """ - pinned_tier: Final = self._tier_for_model(pinned_model) + pinned_tier: Final = tier if tier is not None else self._tier_for_model(pinned_model) if pinned_tier is None: return None escalated_tier: Final = self._escalate_tier(pinned_tier) if escalated_tier == pinned_tier: - return pinned_model - return self.get_model_for_tier(escalated_tier) + return _SessionAffinityPin(pinned_model, pinned_tier) + return _SessionAffinityPin( + self.get_model_for_tier(escalated_tier), _built_in_tier_or_none(_tier_name(escalated_tier)) + ) def _vision_verdicts(self, model_name: str) -> tuple[bool | None, ...]: """Declared vision support per deployment serving the name: True, False, or None when @@ -2907,6 +2987,7 @@ class ComplexityRouter(CustomLogger): resolved_messages: Sequence[Mapping[str, object]] | None, request_kwargs: dict, # mutable-ok: same shape the hook receives context_fit: _RequestContextFit | None = None, + retained_pin: _SessionAffinityPin | None = None, ) -> PreRoutingHookResponse: """Replace a routed model that cannot accept this request's image input. @@ -2955,6 +3036,7 @@ class ComplexityRouter(CustomLogger): repick_messages, # pyright: ignore[reportArgumentType] # hook-resolved message dicts; the pick only reads them request_kwargs, allowed_models=tuple(entry for entry in pools.get(capable, ()) if entry in eligible), + retained_pin=retained_pin, ) elif self._modality_default_model_usable(request_kwargs, resolved_messages, eligible): new_tier = None @@ -3098,6 +3180,7 @@ class ComplexityRouter(CustomLogger): resolved_messages: Sequence[Mapping[str, object]] | None, request_kwargs: dict, # mutable-ok: same shape the hook receives context_fit: _RequestContextFit | None = None, + retained_pin: _SessionAffinityPin | None = None, ) -> PreRoutingHookResponse: """Try compatible tier recovery before the default, preserving request policy and fit.""" decision: Final = response.routing_decision @@ -3155,6 +3238,7 @@ class ComplexityRouter(CustomLogger): repick_messages, # pyright: ignore[reportArgumentType] # hook-resolved message dicts; the pick only reads them request_kwargs, allowed_models=live, + retained_pin=retained_pin, ) except ValueError as exc: verbose_router_logger.debug( @@ -3247,8 +3331,13 @@ class ComplexityRouter(CustomLogger): """The adaptive feedback loop reads its chosen-model marker from request metadata; a gate rewrite must move the marker with the model or rewards land on the displaced one.""" metadata: Final = request_kwargs.get("metadata") - if isinstance(metadata, dict) and metadata.get("adaptive_router_chosen_model") == old_model: + if not isinstance(metadata, dict): + return + if metadata.get("adaptive_router_chosen_model") == old_model: metadata["adaptive_router_chosen_model"] = new_model + decision: Final = metadata.get("adaptive_router_decision") + if isinstance(decision, dict) and decision.get("chosen_model") == old_model: + decision["chosen_model"] = new_model def _lexical_tier_override(self, user_message: str) -> KeywordOverride | None: """When keyword_tier_rules match literally, the most-severe matched tier wins. @@ -3561,25 +3650,42 @@ class ComplexityRouter(CustomLogger): if cache_key is not None and pin_replay_allowed: pinned_value: Final = await self.litellm_router_instance.cache.async_get_cache(key=cache_key) - pinned_pin: Final = _parse_session_affinity_pin(pinned_value) + pinned_pin: Final = _parse_session_affinity_pin(pinned_value, self.config.tier_names()) if pinned_pin is not None: - routed_model: str | None = pinned_pin.model - pin_escalation_keyword: str | None = None - if self.escalation_keywords: - user_message: Final = ( - _newest_turn_ask(resolved_messages, marker_pairs) if resolved_messages else None + user_message: Final = _newest_turn_ask(resolved_messages, marker_pairs) if resolved_messages else None + pin_escalation_keyword: Final = ( + self._matched_escalation_keyword(user_message) if user_message is not None else None + ) + selected_pin: Final = ( + self._escalated_pin(pinned_pin.model, pinned_pin.tier) + if pin_escalation_keyword is not None + else _SessionAffinityPin( + pinned_pin.model, + pinned_pin.tier if pinned_pin.tier is not None else self._tier_for_model(pinned_pin.model), ) - if user_message is not None: - pin_escalation_keyword = self._matched_escalation_keyword(user_message) - if pin_escalation_keyword is not None: - routed_model = self._escalated_pin(pinned_pin.model) - if routed_model is not None: - escalated: Final = routed_model != pinned_pin.model - resolved_pin_tier: Final = ( - pinned_pin.tier - if not escalated and pinned_pin.tier is not None - else self._tier_for_model(routed_model) + ) + if selected_pin is not None: + escalated: Final = selected_pin.model != pinned_pin.model or ( + pin_escalation_keyword is not None + and pinned_pin.tier is not None + and selected_pin.tier != pinned_pin.tier ) + resolved_pin_tier: Final = selected_pin.tier + session_model: Final = ( + await self._pin_model_for_tier( + resolved_pin_tier, + selected_pin.model, + tuple(self._tier_pools().get(_tier_name(resolved_pin_tier), ())), + request_kwargs, + ) + if escalated and resolved_pin_tier is not None + else selected_pin.model + ) + retained_pin: Final = _SessionAffinityPin(session_model, resolved_pin_tier) + if resolved_pin_tier is not None: + await self._pin_model_for_tier( + resolved_pin_tier, session_model, (session_model,), request_kwargs + ) # The floor outranks the pin because plan mode is a transient state of the # session, not a request to move it: the turns carrying the sentinel route at # the floor, and the stored pin deliberately keeps the session's own model so @@ -3590,16 +3696,28 @@ class ComplexityRouter(CustomLogger): plan_floored: Final = ( pinned_tier is not None and self._apply_plan_mode_floor(pinned_tier) != pinned_tier ) - session_model: Final = routed_model - if plan_floored and pinned_tier is not None: - routed_model = self.get_model_for_tier(self._apply_plan_mode_floor(pinned_tier)) - pin_source_tier: Final = self._tier_for_model(routed_model) + floor_model: Final = ( + await self._pick_model_for_tier( + self._apply_plan_mode_floor(pinned_tier), + messages, + resolved_messages, + request_kwargs, + retained_pin=retained_pin, + ) + if plan_floored and pinned_tier is not None + else session_model + ) + pin_source_tier: Final = ( + self._apply_plan_mode_floor(pinned_tier) + if plan_floored and pinned_tier is not None + else resolved_pin_tier + ) pin_placement: Final = ( await self._context_window_placement( pin_source_tier, resolved_messages, request_kwargs, - pool_override=(routed_model,), + pool_override=(floor_model,), context_fit=context_fit, ) if pin_source_tier is not None @@ -3612,11 +3730,18 @@ class ComplexityRouter(CustomLogger): and _tier_name(pin_placement.tier) != _tier_name(pin_source_tier) else None ) - if pin_placement is not None and pin_context_original_tier is not None: - # The stored pin below keeps the session's own model on purpose. - routed_model = self._pick_from_tier_value( - pin_placement.allowed_models, _tier_name(pin_placement.tier) + routed_model: Final = ( + await self._pick_model_for_tier( + pin_placement.tier, + messages, + resolved_messages, + request_kwargs, + allowed_models=pin_placement.allowed_models, + retained_pin=retained_pin, ) + if pin_placement is not None and pin_context_original_tier is not None + else floor_model + ) # Refresh the TTL on every hit so an active session doesn't lose its # pin mid-conversation just because it outlives the original write. await self.litellm_router_instance.cache.async_set_cache( @@ -3644,7 +3769,7 @@ class ComplexityRouter(CustomLogger): routed_pin_tier: Final = ( pin_placement.tier if pin_placement is not None and pin_context_original_tier is not None - else (self._tier_for_model(routed_model) if plan_floored else resolved_pin_tier) + else pin_source_tier ) session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model) has_original_messages: Final = messages is not None and len(messages) > 0 @@ -3671,12 +3796,14 @@ class ComplexityRouter(CustomLogger): resolved_messages, request_kwargs, context_fit, + retained_pin, ), messages, input, resolved_messages, request_kwargs, context_fit, + retained_pin, ) ) @@ -3961,13 +4088,21 @@ class ComplexityRouter(CustomLogger): housekeeping_ceiling: Final = tier if outcome.cause == "housekeeping" else None # A context-escalated tier becomes the hard floor: a floor the bandit can slide # under is not a floor. - routed_model = self._soft_floor_pick( + adaptive_floor: Final = tier if context_original_tier is not None else plan_floor + adaptive_fit: Final = context_placement.holdable_models if context_placement is not None else None + sampled_model: Final = self._soft_floor_pick( tier, ask, request_kwargs, - hard_floor=tier if context_original_tier is not None else plan_floor, + hard_floor=adaptive_floor, hard_ceiling=housekeeping_ceiling, - fit_filter=context_placement.holdable_models if context_placement is not None else None, + fit_filter=adaptive_fit, + ) + routed_model = await self._pin_model_for_tier( # rebind-ok: reuse the eligible tier winner + tier, + sampled_model, + self._adaptive_candidate_models(tier, adaptive_floor, housekeeping_ceiling, adaptive_fit), + request_kwargs, ) adaptive: Final = self._ensure_adaptive_router() if adaptive is not None: diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 1f1b5a5cc4b..bfc83f8dcab 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1256,20 +1256,16 @@ class ComplexityRouterConfig(BaseModel): deployment_affinity: bool = Field( default=True, description=( - "When True and a session_id is resolvable on the request, pin the deployment chosen " - "inside each routed model group and reuse it whenever the session returns to that " - "group, without pinning which group the session routes to. Independent of " - "session_affinity, which pins the model group instead (and always carries this " - "deployment pin with it): with session_affinity off, " - "every turn is still classified on its own merits while a session that escalates to a " - "stronger tier and comes back still lands on the deployment it used before, which is " - "what keeps a provider prompt cache warm. Pins are held per model group, so switching " - "tiers does not disturb the pin left behind in the previous group. On by default " - "because re-shuffling a conversation across deployments of the same model discards " - "that cache for no benefit; set False to keep every turn load-balanced across the " - "group, which is what a deployment set with tight per-deployment rate limits wants. " - "Inert when no session_id is resolvable, since there is nothing to key a pin on, and " - "suppressed when plugins are configured, for the same reason session_affinity is." + "When True and a client session_id is resolvable, reuse the session's chosen model " + "for each classified tier and its deployment within each model group. With " + "session_affinity off, every turn is still classified: moving to another tier leaves " + "the previous tier's model pin intact for a later return. Pins yield to current " + "candidate, context, modality, and availability constraints. Adaptive selection chooses " + "the initial model from its eligible pool, then reuses that choice per tier. This " + "reduces avoidable provider prompt-cache misses; it does not guarantee cache hits. " + "Set False to select models and load-balance deployments on every turn, unless " + "session_affinity or user_turn classification requires a pin. Inert without a client " + "session_id and suppressed when plugins are configured." ), ) session_affinity_ttl_seconds: int = Field( @@ -1277,7 +1273,7 @@ class ComplexityRouterConfig(BaseModel): gt=0, description=( "TTL for the session affinity pin; refreshed on every cache hit. Bounds both the " - "session_affinity model pin and the deployment_affinity deployment pin, so it measures " + "session_affinity model pin and the deployment_affinity per-tier model and deployment pins, so it measures " "idle time for the session's routing decisions rather than total session length" ), ) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index c7eb46046ef..3b88ac2eb00 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -13,13 +13,13 @@ where routing to a consistent deployment is still beneficial. """ import hashlib -import json from collections.abc import Mapping, Sequence from typing import Any, Final, cast -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_router_logger +from litellm.caching.affinity_cache import claim_affinity_pin, claim_affinity_pin_in_memory, set_local_affinity_pin from litellm.caching.dual_cache import DualCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span @@ -28,8 +28,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes -class DeploymentAffinityCacheValue(TypedDict): - model_id: str +class DeploymentAffinityCacheValue(TypedDict, closed=True): + model_id: ReadOnly[str] VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( @@ -60,19 +60,6 @@ def warn_on_unknown_model_group_affinity_flags(model_group_affinity_config: Mapp ) -_CLAIM_PIN_SCRIPT: Final = """ -local current = redis.call('GET', KEYS[1]) -if current == false then - redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) - return ARGV[1] -end -if current == ARGV[1] then - redis.call('EXPIRE', KEYS[1], ARGV[2]) -end -return current -""" - - class DeploymentAffinityCheck(CustomLogger): """ Router deployment affinity callback. @@ -255,34 +242,33 @@ class DeploymentAffinityCheck(CustomLogger): return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}" @staticmethod - def _get_session_id_from_metadata_dict(metadata: dict) -> str | None: + def _get_session_id_from_metadata_dict(metadata: Mapping[object, object]) -> str | None: session_id: Final = metadata.get("session_id") if session_id is None or metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return None return str(session_id) @staticmethod - def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: + def _iter_metadata_dicts(request_kwargs: Mapping[str, object]) -> tuple[Mapping[object, object], ...]: """ Return all metadata dicts available on the request. Depending on the endpoint, Router may populate `metadata` or `litellm_metadata`. Users may also send one or both, so we check both (rather than using `or`). """ - metadata_dicts: Final[list[dict]] = [] - for key in ("litellm_metadata", "metadata"): - md = request_kwargs.get(key) - if isinstance(md, dict): - metadata_dicts.append(md) - return metadata_dicts + return tuple( + cast(Mapping[object, object], metadata) # cast-ok: isinstance proves mapping shape; values remain opaque + for key in ("litellm_metadata", "metadata") + if isinstance(metadata := request_kwargs.get(key), dict) + ) @staticmethod - def _first_metadata_value(metadata_dicts: Sequence[dict], key: str) -> str | None: + def _first_metadata_value(metadata_dicts: Sequence[Mapping[object, object]], key: str) -> str | None: value: Final = next((metadata[key] for metadata in metadata_dicts if metadata.get(key) is not None), None) return None if value is None else str(value) @classmethod - def _get_user_key_from_request_kwargs(cls, request_kwargs: dict) -> str | None: + def get_user_key_from_request_kwargs(cls, request_kwargs: Mapping[str, object]) -> str | None: """ Extract a stable affinity key from request kwargs. @@ -334,74 +320,17 @@ class DeploymentAffinityCheck(CustomLogger): return None def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: - """The one owner of authoritative local pin writes: a plain set keeps a live - key's original expiry (`allow_ttl_override`), so the entry is replaced to make - the TTL real. Every local pin write goes through here so the redis-winner sync - and the pod-local claim can never disagree about expiry again.""" - self.cache.in_memory_cache.delete_cache(cache_key) - self.cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds) + set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds) async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: - """First-writer-wins pin write: store `pin_value` only when the key is absent and - return the deployment id the key holds afterwards, so a caller learns whether it won - by comparing against its own id, and None when the stored value is one no reader can - interpret. Concurrent claimers converge on the - first write instead of the last. Re-claiming with the stored value refreshes its - TTL, the same keepalive the complexity router's model pin documents: an active - session must not lose its pin mid-conversation just because it outlives the - original write, so the affinity TTL (the Router's - `deployment_affinity_ttl_seconds`, or a pre-routing hook's per-request - `session_affinity_ttl_seconds` override) bounds idle time, not total - session length. On Redis one Lua script does the get-or-set-or-refresh - atomically (same registration seam the rate limiters use) and the in-memory - tier is synchronized to the winner; without Redis, and whenever Redis is - unreachable, the pod-local check-and-set below stands in and is atomic because it - runs synchronously on the event loop. Degrading to a pod-local claim rather than - propagating the fault is what keeps same-pod stickiness through a Redis blip: the - caller only logs this result, so an escaping error would leave the session with no - pin at all and reshuffle every turn for the outage, which is worse than losing - cross-pod agreement. The redis tier is - resolved per call because the proxy attaches it after Router construction - (`Router._update_redis_cache`); the compiled script is cached per event loop - underneath the registration seam. - """ - redis_cache: Final = self.cache.redis_cache - if redis_cache is not None: - try: - claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) - raw: Final = await claim_script(keys=(cache_key,), args=(json.dumps(pin_value), int(ttl_seconds))) - decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw - if not isinstance(decoded, str): - return pin_value["model_id"] - try: - winner: object = json.loads(decoded) - except json.JSONDecodeError: - winner = decoded - self._set_local_pin(cache_key=cache_key, value=winner, ttl_seconds=ttl_seconds) - return self._pinned_model_id(winner) - except Exception as e: # noqa: BLE001 # any Redis/Lua failure degrades to the pod-local claim, never unpins - verbose_router_logger.debug( - "DeploymentAffinityCheck: redis pin claim failed, falling back to pod-local claim. error=%s", e - ) - - return self._claim_pin_in_memory(cache_key=cache_key, pin_value=pin_value, ttl_seconds=ttl_seconds) + winner: Final = await claim_affinity_pin(self.cache, cache_key, pin_value, ttl_seconds) + return self._pinned_model_id(winner) def _claim_pin_in_memory( self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int ) -> str | None: - """Pod-local half of the claim, used when no Redis tier is attached and as the - fallback when the Redis claim fails. Mirrors the Lua script exactly, including - the keepalive: re-claiming with the stored value slides the idle window through - `_set_local_pin`. Both branches stay synchronous, hence atomic on the event - loop.""" - existing: Final = self.cache.in_memory_cache.get_cache(cache_key) - if existing is not None: - existing_model_id: Final = self._pinned_model_id(existing) - if existing_model_id == pin_value["model_id"]: - self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds) - return existing_model_id - self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds) - return pin_value["model_id"] + winner: Final = claim_affinity_pin_in_memory(self.cache, cache_key, pin_value, ttl_seconds) + return self._pinned_model_id(winner) @staticmethod def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: @@ -465,7 +394,7 @@ class DeploymentAffinityCheck(CustomLogger): enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None ) user_key: Final = ( - self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs) + self.get_user_key_from_request_kwargs(request_kwargs=request_kwargs) if (session_affinity_active or enable_user_key) else None ) @@ -580,7 +509,7 @@ class DeploymentAffinityCheck(CustomLogger): return None user_key: Final = ( - self._get_user_key_from_request_kwargs(request_kwargs=kwargs) + self.get_user_key_from_request_kwargs(request_kwargs=kwargs) if (enable_user_key or session_affinity_active) else None ) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index dba44d1e2e8..cd72388ea21 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -9,7 +9,7 @@ import json import logging import sys import time -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from copy import deepcopy from functools import partial from typing import Dict, Final, List, Literal @@ -31,6 +31,7 @@ from litellm.router_utils.auto_router_model_naming import ( ) from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( OUTPUT_TOKEN_CEILING_PARAMS, RETURN_RAW_MODEL_NAME_METADATA_KEY, @@ -70,6 +71,7 @@ from litellm.router_strategy.complexity_router.tier_predictor import ( from litellm.types.router import ( Deployment, LiteLLM_Params, + PreRoutingHookResponse, RouterErrors, TaggedPreRoutingStrategy, ) @@ -5574,6 +5576,387 @@ class TestRoutingDecisionCauseLogging: assert "cause=semantic_keyword_match" not in router_log_capture.text +class TestTierModelAffinity: + @staticmethod + async def _route( + router: ComplexityRouter, + metadata: Mapping[str, object], + proposed_model: str, + prompt: str = "compact", + messages: list[dict[str, object]] | None = None, + ) -> PreRoutingHookResponse: + def choose(candidates: Sequence[str]) -> str: + return proposed_model if proposed_model in candidates else candidates[0] + + request_metadata: Final = dict(metadata) + with patch( # test-quality-ok: [TQ008] alternate proposals make affinity reuse deterministic + "litellm.router_strategy.complexity_router.complexity_router.random.choice", + side_effect=choose, + ): + result: Final = await router.async_pre_routing_hook( + model="affinity-router", + request_kwargs={"metadata": request_metadata}, + messages=messages if messages is not None else [{"role": "user", "content": prompt}], + ) + assert result is not None + if router.config.adaptive: + assert request_metadata["adaptive_router_chosen_model"] == result.model + return result + + @staticmethod + def _router( + mock_router_instance: MagicMock, + adaptive: bool = False, + deployment_affinity: bool = True, + plugins: bool = False, + ) -> ComplexityRouter: + mock_router_instance.cache = DualCache() + mock_router_instance.model_list = [] + mock_router_instance.model_name_to_deployment_indices = {} + return ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": { + tier: [ + {"model_name": model, "litellm_params": {"temperature": temperature}} + for model in ("model-a", "model-b") + ] + for tier, temperature in (("SIMPLE", 0.1), ("REASONING", 0.9)) + }, + "adaptive": adaptive, + "deployment_affinity": deployment_affinity, + "session_affinity": False, + **({"plugins": [_DummyPlugin()]} if plugins else {}), + }, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("adaptive", [False, True]) + async def test_reuses_model_per_tier_without_pinning_classification( + self, mock_router_instance: MagicMock, adaptive: bool + ) -> None: + router: Final = self._router(mock_router_instance, adaptive=adaptive) + metadata: Final = {"session_id": "same-session"} + first: Final = await self._route(router, metadata, "model-a") + if adaptive: + from litellm.router_strategy.adaptive_router.bandit import BanditCell + from litellm.router_strategy.adaptive_router.classifier import classify_prompt + + bandit: Final = router._ensure_adaptive_router() + assert bandit is not None + bandit._cells[(classify_prompt("compact"), "model-a")] = BanditCell(alpha=5.0, beta=5.0) + repeated: Final = await self._route(router, metadata, "model-b") + reasoning: Final = await self._route( + router, metadata, "model-b", "Let's think step by step and reason through this problem carefully." + ) + returned: Final = await self._route(router, metadata, "model-b") + + assert (first.model, repeated.model, reasoning.model, returned.model) == ( + "model-a", "model-a", "model-b", "model-a" + ) + assert tuple(result.routing_decision["tier"] for result in (first, repeated, reasoning, returned)) == ( + "SIMPLE", "SIMPLE", "REASONING", "SIMPLE" + ) + assert returned.litellm_params == {"temperature": 0.1} + assert reasoning.litellm_params == {"temperature": 0.9} + + @pytest.mark.asyncio + @pytest.mark.parametrize("identity_key", ["user_api_key_hash", "user_api_key_user_id"]) + async def test_isolates_sessions_and_authenticated_callers( + self, mock_router_instance: MagicMock, identity_key: str + ) -> None: + router: Final = self._router(mock_router_instance) + first_caller: Final = {"session_id": "shared", identity_key: "caller-a"} + other_caller: Final = {"session_id": "shared", identity_key: "caller-b"} + other_session: Final = {"session_id": "separate", identity_key: "caller-a"} + + assert (await self._route(router, first_caller, "model-a")).model == "model-a" + assert (await self._route(router, other_caller, "model-b")).model == "model-b" + assert (await self._route(router, other_session, "model-b")).model == "model-b" + assert (await self._route(router, first_caller, "model-b")).model == "model-a" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "metadata,deployment_affinity,plugins", + [ + ({}, True, False), + ({"session_id": "generated", SESSION_ID_GENERATED_METADATA_KEY: True}, True, False), + ({"session_id": "provided"}, False, False), + ({"session_id": "provided"}, True, True), + ], + ids=["absent-session", "generated-session", "disabled", "plugin-policy"], + ) + async def test_does_not_pin_without_eligible_session( + self, + mock_router_instance: MagicMock, + metadata: Mapping[str, object], + deployment_affinity: bool, + plugins: bool, + ) -> None: + router: Final = self._router( + mock_router_instance, deployment_affinity=deployment_affinity, plugins=plugins + ) + assert (await self._route(router, metadata, "model-a")).model == "model-a" + assert (await self._route(router, metadata, "model-b")).model == "model-b" + + @pytest.mark.asyncio + @pytest.mark.parametrize("adaptive", [False, True]) + async def test_replaces_pin_outside_the_context_candidate_domain(self, adaptive: bool) -> None: + router: Final = ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=_windowed_router(_SMALL, _BIG), + complexity_router_config={ + "tiers": {"SIMPLE": ["small-model", "big-model"]}, + "adaptive": adaptive, + "deployment_affinity": True, + "session_affinity": False, + }, + ) + metadata: Final = {"session_id": "growing-context"} + assert (await self._route(router, metadata, "small-model")).model == "small-model" + oversized: Final = await router.async_pre_routing_hook( + model="affinity-router", + request_kwargs={"metadata": dict(metadata)}, + messages=_OVERSIZED_TURNS, + ) + assert oversized is not None + assert oversized.model == "big-model" + assert oversized.routing_decision["tier"] == "SIMPLE" + assert (await self._route(router, metadata, "small-model")).model == "big-model" + + @pytest.mark.asyncio + @pytest.mark.parametrize("session_affinity", [False, True], ids=["user-turn", "session-affinity"]) + @pytest.mark.parametrize("gate", ["image", "health"]) + async def test_temporary_replay_gate_keeps_the_held_tiers_model_preference( + self, mock_router_instance: MagicMock, session_affinity: bool, gate: Literal["image", "health"] + ) -> None: + async def get_healthy_deployments( + model: str, + request_kwargs: Mapping[str, object], + messages: Sequence[Mapping[str, object]] | None = None, + input: object = None, + parent_otel_span: object = None, + health_check_probe: bool = False, + ) -> list[dict[str, object]]: + unavailable: Final = ( + gate == "health" + and model == "model-a" + and messages is not None + and bool(messages) + and messages[-1].get("role") == "tool" + ) + return [] if unavailable else [{"model_name": model, "model_info": {"id": f"deployment-{model}"}}] + + cache: Final = DualCache() + mock_router_instance.cache = cache + mock_router_instance.async_get_healthy_deployments = get_healthy_deployments + router: Final = TestModalityRouting._router( + mock_router_instance, + { + "tiers": {"SIMPLE": ["model-a", "model-b"]}, + "deployment_affinity": True, + "session_affinity": session_affinity, + "classification_mode": "every_request" if session_affinity else "user_turn", + "modality_routing": True, + "modality_pin_override": True, + }, + {"model-a": False, "model-b": True}, + ) + metadata: Final = {"session_id": "replay-session"} + continuation: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "compact"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": [IMG_PART] if gate == "image" else "done"}, + ] + assert (await self._route(router, metadata, "model-a")).model == "model-a" + + replayed: Final = await self._route(router, metadata, "model-b", messages=continuation) + assert replayed.model == "model-b" + assert replayed.routing_decision["tier"] == "SIMPLE" + assert replayed.routing_decision["cause"] == ( + "health_failover" + if gate == "health" + else ("modality_pin_override" if session_affinity else "user_turn_continuation") + ) + cache_key: Final = router._get_session_affinity_cache_key("replay-session", {"metadata": metadata}) + assert await cache.async_get_cache(cache_key) == {"model": "model-a", "tier": "SIMPLE"} + + next_ask: Final = await self._route(router, metadata, "model-b") + assert next_ask.model == "model-a" + assert next_ask.routing_decision["tier"] == "SIMPLE" + assert next_ask.routing_decision["cause"] == ( + "session_affinity_pin" if session_affinity else "heuristic_scorer" + ) + + @pytest.mark.asyncio + async def test_user_turn_replay_refreshes_the_model_used_within_its_tier( + self, mock_router_instance: MagicMock + ) -> None: + clock: Final = MagicMock(return_value=100.0) + mock_router_instance.cache = DualCache(in_memory_cache=InMemoryCache(clock=clock)) + router: Final = ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": ["model-a", "model-b"]}, + "classification_mode": "user_turn", + "session_affinity_ttl_seconds": 10, + }, + ) + metadata: Final = {"session_id": "same-session"} + continuation: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "compact"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ] + assert (await self._route(router, metadata, "model-a")).model == "model-a" + clock.return_value = 105.0 + replayed: Final = await self._route(router, metadata, "model-b", messages=continuation) + assert replayed.model == "model-a" + assert replayed.routing_decision["cause"] == "user_turn_continuation" + + clock.return_value = 111.0 + next_ask: Final = await self._route(router, metadata, "model-b") + assert next_ask.model == "model-a" + assert next_ask.routing_decision["tier"] == "SIMPLE" + assert next_ask.routing_decision["cause"] == "heuristic_scorer" + + @pytest.mark.asyncio + async def test_session_escalation_keeps_the_selected_tier_when_models_overlap( + self, mock_router_instance: MagicMock + ) -> None: + cache: Final = DualCache() + mock_router_instance.cache = cache + router: Final = ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": { + "SIMPLE": "base", + **{ + tier: [ + {"model_name": model, "litellm_params": {"temperature": temperature}} + for model in models + ] + for tier, models, temperature in ( + ("MEDIUM", ("shared", "middle"), 0.4), + ("COMPLEX", ("shared", "higher"), 0.8), + ) + }, + }, + "session_affinity": True, + "keyword_tier_rules": [{"keywords": ["visit_complex"], "tier": "COMPLEX"}], + }, + ) + metadata: Final = {"session_id": "same-session"} + assert (await self._route(router, metadata, "higher", "visit_complex")).model == "higher" + cache_key: Final = router._get_session_affinity_cache_key("same-session", {"metadata": metadata}) + await cache.async_set_cache(cache_key, {"model": "base", "tier": "SIMPLE"}, ttl=600) + + result: Final = await self._route(router, metadata, "shared", "LITELLM ESCALATE") + assert result.model == "shared" + assert result.routing_decision["tier"] == "MEDIUM" + assert result.routing_decision["cause"] == "session_affinity_escalation" + assert result.litellm_params == {"temperature": 0.4} + assert await cache.async_get_cache(cache_key) == {"model": "shared", "tier": "MEDIUM"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "stale_tier", + ["NON_REASONING", "REMOVED_TIER", 7, []], + ids=["inactive-tier", "unknown-tier", "integer-tier", "list-tier"], + ) + @pytest.mark.parametrize( + "prompt,expected_model,expected_tier", + [("compact", "model-a", "SIMPLE"), ("LITELLM ESCALATE", "model-b", "MEDIUM")], + ids=["ordinary-replay", "escalation"], + ) + async def test_reclassifies_session_pin_outside_the_active_tier_ladder( + self, + mock_router_instance: MagicMock, + stale_tier: object, + prompt: str, + expected_model: str, + expected_tier: str, + ) -> None: + cache: Final = DualCache() + mock_router_instance.cache = cache + router: Final = ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "model-a", "MEDIUM": "model-b"}, + "session_affinity": True, + }, + ) + metadata: Final = {"session_id": "same-session"} + cache_key: Final = router._get_session_affinity_cache_key("same-session", {"metadata": metadata}) + await cache.async_set_cache(cache_key, {"model": "model-a", "tier": stale_tier}, ttl=600) + + result: Final = await self._route(router, metadata, expected_model, prompt) + + assert result.model == expected_model + assert result.routing_decision["tier"] == expected_tier + assert result.routing_decision["cause"] == "heuristic_scorer" + assert await cache.async_get_cache(cache_key) == {"model": expected_model, "tier": expected_tier} + + @pytest.mark.asyncio + @pytest.mark.parametrize("classification_mode", ["every_request", "user_turn"]) + async def test_custom_tier_keeps_its_own_model( + self, mock_router_instance: MagicMock, classification_mode: Literal["every_request", "user_turn"] + ) -> None: + mock_router_instance.cache = DualCache() + router: Final = ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=_custom_tier_config( + tiers={"SIMPLE": ["model-a", "model-b"], "SECURITY_REVIEW": ["model-a", "model-b"], "COMPLEX": "model-a"}, + deployment_affinity=True, + classification_mode=classification_mode, + keyword_tier_rules=[ + {"keywords": ["compact"], "tier": "SIMPLE"}, + {"keywords": ["audit"], "tier": "SECURITY_REVIEW"}, + ], + ), + ) + metadata: Final = {"session_id": "custom-session"} + assert (await self._route(router, metadata, "model-a")).model == "model-a" + assert (await self._route(router, metadata, "model-b", "audit")).model == "model-b" + assert (await self._route(router, metadata, "model-b")).model == "model-a" + retained: Final = await self._route(router, metadata, "model-a", "audit") + assert retained.model == "model-b" + assert retained.routing_decision["tier"] == "SECURITY_REVIEW" + if classification_mode == "user_turn": + continuation: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "audit"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ] + replayed: Final = await self._route(router, metadata, "model-a", messages=continuation) + assert replayed.model == "model-b" + assert replayed.routing_decision["tier"] == "SECURITY_REVIEW" + assert replayed.routing_decision["cause"] == "user_turn_continuation" + + class TestSessionAffinity: """Test the session_affinity sticky-routing behavior (off by default).""" @@ -5638,11 +6021,8 @@ class TestSessionAffinity: tier_pinned, deployment_pinned, ): - """deployment_affinity pins the deployment inside each routed group without pinning which - group the session routes to, so with session_affinity off the tier must still reclassify - on every turn while the marker the Router stamps is still emitted. Turn 1 classifies - REASONING and turn 2 SIMPLE, so a reclassified turn 2 moves model while a tier-pinned one - does not. plugins suppress both pins, since a stale pin would bypass the plugin pipeline.""" + """Deployment affinity retains a model per tier while classification continues. + Session affinity keeps the first tier too; plugins suppress both affinity policies.""" mock_router_instance.cache = DualCache() router = ComplexityRouter( model_name="test-router", @@ -5692,8 +6072,7 @@ class TestSessionAffinity: @pytest.mark.asyncio async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): - """Regression: session_affinity defaults to False, so a shared session_id must NOT - pin the first turn's model; every turn is classified on its own merits.""" + """With session_affinity off, a shared session can move from REASONING to SIMPLE.""" assert "session_affinity" not in basic_config mock_router_instance.cache = DualCache() router = ComplexityRouter( @@ -5848,7 +6227,7 @@ class TestSessionAffinity: @pytest.mark.asyncio async def test_respects_ttl_seconds(self, mock_router_instance, basic_config): - cache = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value=None) mock_router_instance.cache = cache router = ComplexityRouter( @@ -5872,7 +6251,7 @@ class TestSessionAffinity: async def test_ttl_refreshed_on_cache_hit(self, mock_router_instance, basic_config): """Regression: a pinned turn must refresh the TTL, not just the first write -- otherwise a session outliving session_affinity_ttl_seconds silently loses its pin.""" - cache = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value="o1-preview") mock_router_instance.cache = cache router = ComplexityRouter( @@ -7112,7 +7491,8 @@ class TestEscalationKeywords: complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]}}, ) for pinned in ("o1-a", "o1-b", "o1-c"): - assert router._escalated_pin(pinned) == pinned + escalated: Final = router._escalated_pin(pinned) + assert (escalated.model, escalated.tier) == (pinned, "REASONING") @pytest.mark.asyncio async def test_session_escalation_at_ceiling_keeps_multi_model_pin(self, mock_router_instance): @@ -12009,7 +12389,7 @@ async def test_session_pin_uses_recorded_tier_when_model_is_in_multiple_tiers(mo @pytest.mark.asyncio async def test_session_pin_survives_json_list_round_trip(mock_router_instance): - cache = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value=["shared", "SIMPLE"]) mock_router_instance.cache = cache router = ComplexityRouter( @@ -12988,7 +13368,7 @@ class TestModalityRouting: {"role": "user", "content": [{"type": "text", "text": "quick lookup: what is this?"}, IMG_PART]} ] elif path.startswith(("pin_kept", "pin_replacement", "pin_override")): - cache = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) mock_router_instance.cache = cache config["session_affinity"] = True @@ -13178,7 +13558,7 @@ class TestModalityRouting: @pytest.mark.asyncio async def test_pin_override_serves_the_image_turn_without_repinning(self, mock_router_instance): """The override is for one request: the session keeps the model it was pinned to.""" - cache = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) mock_router_instance.cache = cache router = self._router( @@ -13211,7 +13591,7 @@ class TestModalityRouting: @pytest.mark.asyncio async def test_pin_override_with_no_capable_model_rejects_and_keeps_the_pin(self, mock_router_instance): """The clear 400 replaces the provider's, and a rejected turn must not cost the session its pin.""" - cache = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) mock_router_instance.cache = cache router = self._router( @@ -14322,18 +14702,37 @@ class TestTierHealthFailover: cooling=("id-a1",), raises_for={"exhausted-b": raised}, ) - key = router._get_session_affinity_cache_key("sess-exhausted", {}) - await router.litellm_router_instance.cache.async_set_cache( - key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 - ) - results = [ - await router.async_pre_routing_hook( - model="m", request_kwargs={"metadata": {"session_id": "sess-exhausted"}}, messages=self.SIMPLE_MESSAGE + sessions: Final = tuple(f"sess-exhausted-{sample}" for sample in range(20)) + await asyncio.gather( + *( + router.litellm_router_instance.cache.async_set_cache( + key=router._get_session_affinity_cache_key(session_id, {}), + value={"model": "dead-a", "tier": "SIMPLE"}, + ttl=600, + ) + for session_id in sessions ) - for _ in range(20) + ) + results: Final = [ + await router.async_pre_routing_hook( + model="m", request_kwargs={"metadata": {"session_id": session_id}}, messages=self.SIMPLE_MESSAGE + ) + for session_id in sessions ] assert {r.model for r in results} == expected + def choose_other(candidates: Sequence[str]) -> str: + return next((model for model in candidates if model != results[0].model), candidates[0]) + + with patch( # test-quality-ok: [TQ008] an alternate healthy proposal proves retained affinity across failover + "litellm.router_strategy.complexity_router.complexity_router.random.choice", + side_effect=choose_other, + ): + retained: Final = await router.async_pre_routing_hook( + model="m", request_kwargs={"metadata": {"session_id": sessions[0]}}, messages=self.SIMPLE_MESSAGE + ) + assert retained.model == results[0].model + @pytest.mark.asyncio async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(self, mock_router_instance): """The owner answers an unconfigured group with BadRequestError. Reading that as live diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index cf48888600e..780300bf9e1 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,3 +1,5 @@ +import asyncio +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -6,7 +8,9 @@ import pytest import json import litellm +from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, @@ -558,6 +562,124 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): assert second == "our-deployment" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stored", "expected"), + [ + ({"model": "first"}, {"model": "first"}), + ('{ "model" : "first" }', {"model": "first"}), + ({"model": "removed"}, {"model": "second"}), + ({"model": "first", "extra": "stale"}, {"model": "second"}), + ({"model_id": "first"}, {"model": "second"}), + ("first", {"model": "second"}), + (None, {"model": "second"}), + ], +) +async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl( + stored: object, expected: object +) -> None: + clock: Final = MagicMock(return_value=100.0) + cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) + cache.in_memory_cache.set_cache("tier-pin", stored, ttl=10) + clock.return_value = 105.0 + + winner: Final = await claim_affinity_pin( + cache, "tier-pin", {"model": "second"}, 30, + eligible_values=({"model": "first"}, {"model": "second"}), + ) + + assert winner == expected + assert cache.in_memory_cache.ttl_dict["tier-pin"] == 135.0 + clock.return_value = 111.0 + assert cache.in_memory_cache.get_cache("tier-pin") == expected + clock.return_value = 136.0 + assert cache.in_memory_cache.get_cache("tier-pin") is None + + +@pytest.mark.asyncio +async def test_concurrent_eligible_claims_return_one_winner() -> None: + cache: Final = DualCache() + candidates: Final = ({"model": "first"}, {"model": "second"}) + winners: Final = await asyncio.gather(*( + claim_affinity_pin( + cache, "tier-pin", candidates[index % 2], 30, + eligible_values=candidates, + ) + for index in range(20) + )) + + assert winners == [{"model": "first"}] * 20 + assert cache.in_memory_cache.get_cache("tier-pin") == {"model": "first"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stored", "expected", "refresh"), + [ + ({"model_id": 7}, "7", True), + ({"model_id": "other"}, "other", False), + ({"model": "7"}, None, False), + (["7"], None, False), + ], +) +async def test_legacy_deployment_claim_retains_decoder_and_keepalive( + stored: object, expected: str | None, refresh: bool +) -> None: + clock: Final = MagicMock(return_value=100.0) + cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) + callback: Final = DeploymentAffinityCheck( + cache=cache, ttl_seconds=30, + enable_user_key_affinity=False, enable_responses_api_affinity=False, + ) + cache.in_memory_cache.set_cache("deployment-pin", stored, ttl=10) + clock.return_value = 105.0 + + winner: Final = await callback._claim_pin( + "deployment-pin", {"model_id": "7"}, 30 + ) + + assert winner == expected + assert cache.in_memory_cache.ttl_dict["deployment-pin"] == ( + 135.0 if refresh else 110.0 + ) + assert cache.in_memory_cache.get_cache("deployment-pin") == ( + {"model_id": "7"} if refresh else stored + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("raw", "expected", "stored"), + [ + (b'{"model_id": "winner"}', "winner", {"model_id": "winner"}), + ('"winner"', "winner", "winner"), + ("winner", "winner", "winner"), + (b"winner", "winner", "winner"), + ('{"model": "winner"}', None, {"model": "winner"}), + (None, "candidate", None), + (123, "candidate", None), + ({"model_id": "winner"}, "candidate", None), + ], +) +async def test_redis_deployment_claim_preserves_legacy_result_decoding( + raw: object, expected: str | None, stored: object +) -> None: + redis: Final = MagicMock() + redis.async_register_script.return_value = AsyncMock(return_value=raw) + cache: Final = DualCache(redis_cache=redis) + callback: Final = DeploymentAffinityCheck( + cache=cache, ttl_seconds=30, + enable_user_key_affinity=False, enable_responses_api_affinity=False, + ) + + winner: Final = await callback._claim_pin( + "deployment-pin", {"model_id": "candidate"}, 30 + ) + + assert winner == expected + assert cache.in_memory_cache.get_cache("deployment-pin") == stored + + @pytest.mark.asyncio async def test_marker_session_affinity_read_and_write_agree_for_wildcard_groups(): """Wildcard deployments keep the literal pattern as model_name on both the read diff --git a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx index 9022d424369..325362ea177 100644 --- a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx @@ -28,13 +28,13 @@ export const AffinityControls: React.FC<{ onChange({ ...value, deployment_affinity: deploymentAffinity })} - aria-label="Pin a session to one deployment per model group" + aria-label="Pin one model deployment per tier" /> - Pin a session to one deployment per model group + Pin one model deployment per tier - Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to - load-balance every turn. + Reuses the model chosen for each tier and its deployment when available. Requests can still move between tiers. + Turn off to select models and load-balance deployments every turn.