From 833670f7dbe32331f8689171bc9c0310ea11dd0a Mon Sep 17 00:00:00 2001 From: elinacse Date: Sun, 2 Aug 2026 12:20:46 +0530 Subject: [PATCH 01/16] fix(batch): track cost for managed batches with no attributable key/user/team LiteLLM_ManagedObjectTable only stores created_by (user_id) and team_id, never the raw API key hash. A batch created with the master key or a team-less key has both null, so CheckBatchCost's synthetic logging_obj for the completed batch carried no attributable key/user/team/end-user. _should_track_cost_callback silently skipped the DB write in that case (by design, to avoid tracking truly anonymous requests), with no error or warning: batch_processed still became true, but no LiteLLM_SpendLogs row was ever written despite real, already-incurred provider cost. Extend the same allowance already made for unauthenticated pass-through requests to aretrieve_batch's cost event, and pass job.team_id through so a batch's team gets real attribution when one exists. --- .../proxy/common_utils/check_batch_cost.py | 1 + .../proxy/hooks/proxy_track_cost_callback.py | 12 +- .../proxy_unit_tests/test_check_batch_cost.py | 128 ++++++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 17 ++- 4 files changed, 154 insertions(+), 4 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 22f9f40ecd8..0214c6cceb6 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -502,6 +502,7 @@ class CheckBatchCost: }, "metadata": { "user_api_key_user_id": creator_user_id, + "user_api_key_team_id": getattr(job, "team_id", None), **user_info, }, }, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 857429fa89f..2ff7808868d 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -34,11 +34,17 @@ from litellm.types.utils import ( ) from litellm.utils import get_end_user_id_for_cost_tracking -_PASS_THROUGH_CALL_TYPES: frozenset[str] = frozenset( +_UNATTRIBUTED_TRACKABLE_CALL_TYPES: frozenset[str] = frozenset( { CallTypes.pass_through.value, CallTypes.llm_passthrough_route.value, CallTypes.allm_passthrough_route.value, + # CheckBatchCost's synthetic logging_obj for a completed managed batch only ever + # carries user_api_key_user_id (from LiteLLM_ManagedObjectTable.created_by) and + # user_api_key_team_id (from .team_id) -- both are None for batches created with + # the master key or a team-less key, since the table never stores the raw key + # hash. The batch already incurred real provider cost, so track it regardless. + CallTypes.aretrieve_batch.value, } ) @@ -434,6 +440,8 @@ def _should_track_cost_callback( the request with no key/user/team/end-user to attribute spend to. Those requests still forward real provider traffic that operators expect to see in request/usage logs, so they are tracked even when unauthenticated. + The same reasoning applies to a completed managed batch's cost event + (see _UNATTRIBUTED_TRACKABLE_CALL_TYPES). """ # don't run track cost callback if user opted into disabling spend @@ -442,7 +450,7 @@ def _should_track_cost_callback( if user_api_key is not None or user_id is not None or team_id is not None or end_user_id is not None: return True - return call_type in _PASS_THROUGH_CALL_TYPES + return call_type in _UNATTRIBUTED_TRACKABLE_CALL_TYPES def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index a15abd023d8..42499f2ac55 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -420,6 +420,134 @@ class TestCheckBatchCost: ), "update() must include batch_processed=True when column is present" assert update_data["status"] == "complete" + @pytest.mark.asyncio + async def test_completed_batch_with_no_attributable_owner_still_writes_spend_log( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """Regression: a batch created with the master key or a team-less key has + created_by=None and team_id=None on LiteLLM_ManagedObjectTable (the table + never stores the raw key hash). CheckBatchCost's synthetic logging_obj for + such a batch then carries no attributable key/user/team/end-user, and + before the fix _should_track_cost_callback silently skipped the DB write + with no error or warning: batch_processed still became True, but no + LiteLLM_SpendLogs row was ever written. + + Unlike the other tests in this file, this one does NOT mock + litellm_logging.Logging or async_success_handler -- it runs the real + logging pipeline through to _ProxyDBLogger, which is the exact gap that + let the original bug ship undetected. + """ + import litellm + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + mock_job = MagicMock() + mock_job.id = "job-unattributed-1" + mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.created_by = None + mock_job.team_id = None + + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + + # A real LiteLLMBatch (not a bare MagicMock): this test runs the real + # litellm_logging.Logging pipeline, which type-checks the result via + # isinstance(..., LiteLLMBatch) before it will compute/attach a cost. + from litellm.types.utils import LiteLLMBatch + + mock_response = LiteLLMBatch( + id="batch-1", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-input-123", + object="batch", + status="completed", + output_file_id="file-output-123", + ) + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "openai" + mock_deployment.litellm_params.model = "gpt-4" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + mock_file_content = MagicMock() + mock_file_content.content = b'{"id":"req-1"}' + + decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" + + db_logger = _ProxyDBLogger() + mock_update_database = AsyncMock() + + # Unlike the other tests in this file, this one runs the real + # litellm_logging.Logging pipeline, which calls + # _is_base64_encoded_unified_file_id an extra time (checking result.id + # after it's reset to job.unified_object_id). Key off the argument + # instead of a fixed-length side_effect list so the exact call count + # doesn't matter. + def _fake_is_base64_encoded(file_id): + return decoded_id if file_id == mock_job.unified_object_id else None + + with ( + patch( + "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", + side_effect=_fake_is_base64_encoded, + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", + return_value="model-123", + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", + return_value="batch-456", + ), + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=mock_file_content, + ), + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"id": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=( + 0.01, + {"prompt_tokens": 10, "completion_tokens": 5}, + ["gpt-4"], + ), + ), + patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", + return_value=("gpt-4", "openai", None, None), + ), + patch.object(litellm, "_async_success_callback", [db_logger]), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock( + db_spend_update_writer=MagicMock(update_database=mock_update_database), + slack_alerting_instance=MagicMock(customer_spend_alert=AsyncMock()), + ), + ), + patch("litellm.proxy.proxy_server.increment_spend_counters", AsyncMock()), + patch("litellm.proxy.proxy_server.update_cache", AsyncMock()), + ): + await check_batch_cost_instance.check_batch_cost() + + mock_update_database.assert_awaited_once() + assert mock_update_database.call_args.kwargs["response_cost"] == 0.01 + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "the job must still be marked processed once cost tracking succeeds" + ) + @pytest.mark.asyncio async def test_cost_tracking_failure_leaves_job_unprocessed( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router 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 f289148101a..69f04ce2bbe 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 @@ -1186,6 +1186,7 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): ("pass_through_endpoint", True), ("llm_passthrough_route", True), ("allm_passthrough_route", True), + ("aretrieve_batch", True), ("acompletion", False), ("call_mcp_tool", False), (None, False), @@ -1194,7 +1195,14 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): def test_should_track_cost_callback_pass_through_without_owner(call_type, expected): """Regression for LIT-3782: unauthenticated pass-through requests (auth=false) carry no key/user/team/end-user, yet must still be tracked so they land in - LiteLLM_SpendLogs. Other call types with no owner stay untracked.""" + LiteLLM_SpendLogs. Other call types with no owner stay untracked. + + aretrieve_batch is included for the same reason: CheckBatchCost's synthetic + logging_obj for a completed managed batch only ever carries + user_api_key_user_id/user_api_key_team_id from LiteLLM_ManagedObjectTable, + both of which are None for a batch created with the master key or a + team-less key (the table never stores the raw key hash). Before this fix, + such a batch's cost silently never reached LiteLLM_SpendLogs.""" assert ( _should_track_cost_callback( user_api_key=None, @@ -1211,6 +1219,7 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect "call_type, expect_spend_log", [ ("pass_through_endpoint", True), + ("aretrieve_batch", True), ("acompletion", False), (None, False), ], @@ -1223,7 +1232,11 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( 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 must now be written for pass-through call types while other unauthenticated - calls remain skipped.""" + calls remain skipped. + + aretrieve_batch is included because CheckBatchCost's completed-batch cost + event reaches this same callback with no attributable key/user/team when + the batch was created with the master key or a team-less key.""" logger = _ProxyDBLogger() kwargs = { From 629c228b40e8d1c829c89c62d1cadf58f8a92d38 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 4 Aug 2026 23:56:52 -0700 Subject: [PATCH 02/16] fix(pricing): sync flex/priority tier keys to dated OpenAI snapshot variants Dated snapshots like o4-mini-2025-04-16 were missing the flex and priority cost keys their base alias carries, so service-tier requests against pinned snapshots were billed at standard rates. Sync the tier keys wherever the snapshot's anchor prices match the base alias, and add a drift regression test. --- ...odel_prices_and_context_window_backup.json | 31 ++++++++++++++++++ model_prices_and_context_window.json | 31 ++++++++++++++++++ .../test_litellm/test_model_prices_schema.py | 32 +++++++++++++++++++ 3 files changed, 94 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 02b3cde217a..cfd5bd8dfa0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22176,7 +22176,9 @@ }, "gpt-4.1-2025-04-14": { "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_priority": 8.75e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3.5e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "openai", "max_input_tokens": 1047576, @@ -22184,6 +22186,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_priority": 1.4e-05, "output_cost_per_token_batches": 4e-06, "supported_endpoints": [ "/v1/chat/completions", @@ -22247,7 +22250,9 @@ }, "gpt-4.1-mini-2025-04-14": { "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_priority": 1.75e-07, "input_cost_per_token": 4e-07, + "input_cost_per_token_priority": 7e-07, "input_cost_per_token_batches": 2e-07, "litellm_provider": "openai", "max_input_tokens": 1047576, @@ -22255,6 +22260,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1.6e-06, + "output_cost_per_token_priority": 2.8e-06, "output_cost_per_token_batches": 8e-07, "supported_endpoints": [ "/v1/chat/completions", @@ -22317,7 +22323,9 @@ }, "gpt-4.1-nano-2025-04-14": { "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 1e-07, + "input_cost_per_token_priority": 2e-07, "input_cost_per_token_batches": 5e-08, "litellm_provider": "openai", "max_input_tokens": 1047576, @@ -22325,6 +22333,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-07, + "output_cost_per_token_priority": 8e-07, "output_cost_per_token_batches": 2e-07, "supported_endpoints": [ "/v1/chat/completions", @@ -22393,7 +22402,9 @@ }, "gpt-4o-2024-08-06": { "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_priority": 2.125e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_priority": 4.25e-06, "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22401,6 +22412,7 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_priority": 1.7e-05, "output_cost_per_token_batches": 5e-06, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -22413,7 +22425,9 @@ }, "gpt-4o-2024-11-20": { "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_priority": 2.125e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_priority": 4.25e-06, "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22421,6 +22435,7 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_priority": 1.7e-05, "output_cost_per_token_batches": 5e-06, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -22720,7 +22735,9 @@ }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_priority": 1.25e-07, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_priority": 2.5e-07, "input_cost_per_token_batches": 7.5e-08, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22728,6 +22745,7 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_priority": 1e-06, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { "search_context_size_high": 0.03, @@ -25077,6 +25095,7 @@ "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, + "input_cost_per_token_priority": 2.5e-06, "input_cost_per_token_flex": 2.5e-08, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -29304,13 +29323,19 @@ }, "o3-2025-04-16": { "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 8.75e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_flex": 4e-06, + "output_cost_per_token_priority": 1.4e-05, "supported_endpoints": [ "/v1/responses", "/v1/chat/completions", @@ -29525,13 +29550,19 @@ }, "o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, + "cache_read_input_token_cost_flex": 1.375e-07, + "cache_read_input_token_cost_priority": 5e-07, "input_cost_per_token": 1.1e-06, + "input_cost_per_token_flex": 5.5e-07, + "input_cost_per_token_priority": 2e-06, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_flex": 2.2e-06, + "output_cost_per_token_priority": 8e-06, "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_pdf_input": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bc7330ec99c..b411a5fb483 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -22251,7 +22251,9 @@ }, "gpt-4.1-2025-04-14": { "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_priority": 8.75e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3.5e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "openai", "max_input_tokens": 1047576, @@ -22259,6 +22261,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_priority": 1.4e-05, "output_cost_per_token_batches": 4e-06, "supported_endpoints": [ "/v1/chat/completions", @@ -22322,7 +22325,9 @@ }, "gpt-4.1-mini-2025-04-14": { "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_priority": 1.75e-07, "input_cost_per_token": 4e-07, + "input_cost_per_token_priority": 7e-07, "input_cost_per_token_batches": 2e-07, "litellm_provider": "openai", "max_input_tokens": 1047576, @@ -22330,6 +22335,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 1.6e-06, + "output_cost_per_token_priority": 2.8e-06, "output_cost_per_token_batches": 8e-07, "supported_endpoints": [ "/v1/chat/completions", @@ -22392,7 +22398,9 @@ }, "gpt-4.1-nano-2025-04-14": { "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_priority": 5e-08, "input_cost_per_token": 1e-07, + "input_cost_per_token_priority": 2e-07, "input_cost_per_token_batches": 5e-08, "litellm_provider": "openai", "max_input_tokens": 1047576, @@ -22400,6 +22408,7 @@ "max_tokens": 32768, "mode": "chat", "output_cost_per_token": 4e-07, + "output_cost_per_token_priority": 8e-07, "output_cost_per_token_batches": 2e-07, "supported_endpoints": [ "/v1/chat/completions", @@ -22468,7 +22477,9 @@ }, "gpt-4o-2024-08-06": { "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_priority": 2.125e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_priority": 4.25e-06, "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22476,6 +22487,7 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_priority": 1.7e-05, "output_cost_per_token_batches": 5e-06, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -22488,7 +22500,9 @@ }, "gpt-4o-2024-11-20": { "cache_read_input_token_cost": 1.25e-06, + "cache_read_input_token_cost_priority": 2.125e-06, "input_cost_per_token": 2.5e-06, + "input_cost_per_token_priority": 4.25e-06, "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22496,6 +22510,7 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 1e-05, + "output_cost_per_token_priority": 1.7e-05, "output_cost_per_token_batches": 5e-06, "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -22795,7 +22810,9 @@ }, "gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, + "cache_read_input_token_cost_priority": 1.25e-07, "input_cost_per_token": 1.5e-07, + "input_cost_per_token_priority": 2.5e-07, "input_cost_per_token_batches": 7.5e-08, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22803,6 +22820,7 @@ "max_tokens": 16384, "mode": "chat", "output_cost_per_token": 6e-07, + "output_cost_per_token_priority": 1e-06, "output_cost_per_token_batches": 3e-07, "search_context_cost_per_query": { "search_context_size_high": 0.03, @@ -25152,6 +25170,7 @@ "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, "input_cost_per_token": 5e-08, + "input_cost_per_token_priority": 2.5e-06, "input_cost_per_token_flex": 2.5e-08, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -29379,13 +29398,19 @@ }, "o3-2025-04-16": { "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_flex": 2.5e-07, + "cache_read_input_token_cost_priority": 8.75e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 8e-06, + "output_cost_per_token_flex": 4e-06, + "output_cost_per_token_priority": 1.4e-05, "supported_endpoints": [ "/v1/responses", "/v1/chat/completions", @@ -29600,13 +29625,19 @@ }, "o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, + "cache_read_input_token_cost_flex": 1.375e-07, + "cache_read_input_token_cost_priority": 5e-07, "input_cost_per_token": 1.1e-06, + "input_cost_per_token_flex": 5.5e-07, + "input_cost_per_token_priority": 2e-06, "litellm_provider": "openai", "max_input_tokens": 200000, "max_output_tokens": 100000, "max_tokens": 100000, "mode": "chat", "output_cost_per_token": 4.4e-06, + "output_cost_per_token_flex": 2.2e-06, + "output_cost_per_token_priority": 8e-06, "supports_function_calling": true, "supports_parallel_function_calling": false, "supports_pdf_input": true, diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py index f6f1bf16742..ccb0541d318 100644 --- a/tests/test_litellm/test_model_prices_schema.py +++ b/tests/test_litellm/test_model_prices_schema.py @@ -2,6 +2,7 @@ from __future__ import annotations import importlib.util import json +import re from pathlib import Path import jsonschema @@ -97,3 +98,34 @@ def test_schema_accepts_minimal_and_unknown_optional_fields(committed_schema: di validator = build_validator(committed_schema) assert validator.is_valid({"some-model": {"litellm_provider": "openai"}}) assert validator.is_valid({"some-model": {"litellm_provider": "openai", "brand_new_field": {"nested": True}}}) + + +DATED_VARIANT = re.compile(r"^(.*?)-(\d{4}-\d{2}-\d{2})$") +SERVICE_TIER_SUFFIXES = ("_flex", "_priority") + + +def tier_anchor(tier_key: str) -> str: + matched = next(suffix for suffix in SERVICE_TIER_SUFFIXES if tier_key.endswith(suffix)) + return tier_key[: -len(matched)] + + +def test_dated_variants_carry_base_alias_service_tier_pricing(prices: dict): + drifted = [ + f"{name}: missing {tier_key}={base[tier_key]} (base alias {match.group(1)})" + for name, entry in prices.items() + if isinstance(entry, dict) + for match in [DATED_VARIANT.match(name)] + if match is not None + for base in [prices.get(match.group(1))] + if isinstance(base, dict) + for tier_key in base + if tier_key.endswith(SERVICE_TIER_SUFFIXES) + and tier_anchor(tier_key) in base + and entry.get(tier_anchor(tier_key)) == base[tier_anchor(tier_key)] + and entry.get(tier_key) != base[tier_key] + ] + assert drifted == [], ( + "dated model variants are missing flex/priority pricing their base alias has; " + "sync the tier keys so service-tier requests against pinned snapshots are not " + "billed at standard rates:\n" + "\n".join(drifted) + ) From e0833c4ba361c5873de8f1b3485331bbe5fc7137 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 00:15:52 -0700 Subject: [PATCH 03/16] fix(cost): bill reasoning tokens at the service tier output rate A tier request against a model that publishes tier output pricing but no tier reasoning key (every current Gemini flash entry) billed reasoning tokens at the standard output_cost_per_reasoning_token, undercounting priority and fast traffic where thinking tokens dominate completions generic_cost_per_token now resolves the reasoning rate with explicit precedence: an explicit output_cost_per_reasoning_token_ key wins, then the tier-resolved output rate when the model prices that tier, then the standard reasoning key, then the output base cost. The two tier reasoning keys are wired through ModelInfo so providers can publish real tiered reasoning prices when they exist --- .../litellm_core_utils/llm_cost_calc/utils.py | 24 +++- litellm/types/utils.py | 4 + litellm/utils.py | 4 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 117 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 ++ 5 files changed, 154 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index aeef604510b..5c923e5a8fa 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -681,6 +681,23 @@ def _get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: str | return 1.0 +def _resolve_reasoning_token_cost( + model_info: ModelInfo, + service_tier: str | None, + completion_base_cost: float, +) -> float: + tier_reasoning_key: Final = _get_service_tier_cost_key("output_cost_per_reasoning_token", service_tier) + if model_info.get(tier_reasoning_key) is not None: + tier_reasoning_cost: Final = _get_cost_per_unit(model_info, tier_reasoning_key, None) + if tier_reasoning_cost is not None: + return tier_reasoning_cost + tier_output_key: Final = _get_service_tier_cost_key("output_cost_per_token", service_tier) + if tier_output_key != "output_cost_per_token" and model_info.get(tier_output_key) is not None: + return completion_base_cost + standard_reasoning_cost: Final = _get_cost_per_unit(model_info, "output_cost_per_reasoning_token", None) + return standard_reasoning_cost if standard_reasoning_cost is not None else completion_base_cost + + def generic_cost_per_token( model: str, usage: Usage, @@ -817,9 +834,10 @@ def generic_cost_per_token( ## REASONING COST if not is_text_tokens_total and reasoning_tokens and reasoning_tokens > 0: - _output_cost_per_reasoning_token = _get_cost_per_unit(model_info, "output_cost_per_reasoning_token", None) - _output_cost_per_reasoning_token = ( - _output_cost_per_reasoning_token if _output_cost_per_reasoning_token is not None else completion_base_cost + _output_cost_per_reasoning_token = _resolve_reasoning_token_cost( + model_info=model_info, + service_tier=service_tier, + completion_base_cost=completion_base_cost, ) completion_cost += float(reasoning_tokens) * _output_cost_per_reasoning_token diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 77f83c5b6f8..8e23a380792 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -263,6 +263,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_video_token: Optional[float] # for gemini omni models with video output output_vector_size: Optional[int] output_cost_per_reasoning_token: Optional[float] + output_cost_per_reasoning_token_flex: Optional[float] + output_cost_per_reasoning_token_priority: Optional[float] output_cost_per_video_per_second: Optional[float] # only for vertex ai models output_cost_per_audio_per_second: Optional[float] # only for vertex ai models output_cost_per_second: Optional[float] # for OpenAI Speech models @@ -3308,6 +3310,8 @@ class CustomPricingLiteLLMParams(BaseModel): output_cost_per_image_token: Optional[float] = None output_cost_per_video_token: Optional[float] = None output_cost_per_reasoning_token: Optional[float] = None + output_cost_per_reasoning_token_flex: Optional[float] = None + output_cost_per_reasoning_token_priority: Optional[float] = None output_cost_per_video_per_second: Optional[float] = None output_cost_per_audio_per_second: Optional[float] = None search_context_cost_per_query: Optional[Dict[str, Any]] = None diff --git a/litellm/utils.py b/litellm/utils.py index d24a4dc928f..3dbb35926dd 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5482,6 +5482,10 @@ def _get_model_info_helper( output_cost_per_audio_token=_model_info.get("output_cost_per_audio_token", None), output_cost_per_character=_model_info.get("output_cost_per_character", None), output_cost_per_reasoning_token=_model_info.get("output_cost_per_reasoning_token", None), + output_cost_per_reasoning_token_flex=_model_info.get("output_cost_per_reasoning_token_flex", None), + output_cost_per_reasoning_token_priority=_model_info.get( + "output_cost_per_reasoning_token_priority", None + ), output_cost_per_token_above_128k_tokens=_model_info.get( "output_cost_per_token_above_128k_tokens", None ), 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 9145e5dc76d..f5b9ddfa7d0 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 @@ -2553,3 +2553,120 @@ def test_fast_service_tier_matches_priority_above_the_context_threshold(_local_m assert fast == priority assert fast[0] == pytest.approx(300_000 * 1e-05, rel=1e-9) assert fast[1] == pytest.approx(1_000 * 4.5e-05, rel=1e-9) + + +def test_priority_reasoning_tokens_bill_at_the_priority_output_rate(_local_model_cost_map): + """Regression: gemini-3.5-flash publishes priority output pricing but no priority + reasoning key, so reasoning tokens under priority/fast were billed at the standard + output_cost_per_reasoning_token instead of following the tier's output rate.""" + from litellm.types.utils import Usage + + usage = Usage( + prompt_tokens=1_000, + completion_tokens=5_000, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=4_000), + ) + + model_info = litellm.get_model_info(model="gemini-3.5-flash", custom_llm_provider="gemini") + standard_output_rate = model_info["output_cost_per_token"] + standard_reasoning_rate = model_info["output_cost_per_reasoning_token"] + priority_output_rate = model_info["output_cost_per_token_priority"] + assert priority_output_rate is not None + assert priority_output_rate != standard_reasoning_rate + + standard = generic_cost_per_token( + model="gemini-3.5-flash", usage=usage, custom_llm_provider="gemini", service_tier=None + ) + priority = generic_cost_per_token( + model="gemini-3.5-flash", usage=usage, custom_llm_provider="gemini", service_tier="priority" + ) + fast = generic_cost_per_token( + model="gemini-3.5-flash", usage=usage, custom_llm_provider="gemini", service_tier="fast" + ) + + assert standard[1] == pytest.approx(1_000 * standard_output_rate + 4_000 * standard_reasoning_rate, rel=1e-9) + assert priority[1] == pytest.approx(5_000 * priority_output_rate, rel=1e-9) + assert fast == priority + + +def test_explicit_tier_reasoning_key_wins_over_the_tier_output_rate(): + from litellm.types.utils import Usage + + model_info = { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + "output_cost_per_reasoning_token": 6e-06, + "input_cost_per_token_priority": 2e-06, + "output_cost_per_token_priority": 8e-06, + "output_cost_per_reasoning_token_priority": 1.2e-05, + } + usage = Usage( + prompt_tokens=100, + completion_tokens=1_000, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=600), + ) + + _, completion_cost = generic_cost_per_token( + model="synthetic-model", + usage=usage, + custom_llm_provider="openai", + service_tier="priority", + model_info=model_info, + ) + + assert completion_cost == pytest.approx(400 * 8e-06 + 600 * 1.2e-05, rel=1e-9) + + +def test_null_tier_reasoning_key_falls_back_to_the_tier_output_rate(): + """get_model_info dumps every ModelInfo field, so an unpublished tier reasoning key + arrives as an explicit None and must not shadow the tier output rate.""" + from litellm.types.utils import Usage + + model_info = { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + "output_cost_per_reasoning_token": 6e-06, + "output_cost_per_reasoning_token_priority": None, + "input_cost_per_token_priority": 2e-06, + "output_cost_per_token_priority": 8e-06, + } + usage = Usage( + prompt_tokens=100, + completion_tokens=1_000, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=600), + ) + + _, completion_cost = generic_cost_per_token( + model="synthetic-model", + usage=usage, + custom_llm_provider="openai", + service_tier="priority", + model_info=model_info, + ) + + assert completion_cost == pytest.approx(1_000 * 8e-06, rel=1e-9) + + +def test_tier_request_without_tier_pricing_keeps_the_standard_reasoning_rate(): + from litellm.types.utils import Usage + + model_info = { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 4e-06, + "output_cost_per_reasoning_token": 6e-06, + } + usage = Usage( + prompt_tokens=100, + completion_tokens=1_000, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=600), + ) + + _, completion_cost = generic_cost_per_token( + model="synthetic-model", + usage=usage, + custom_llm_provider="openai", + service_tier="priority", + model_info=model_info, + ) + + assert completion_cost == pytest.approx(400 * 4e-06 + 600 * 6e-06, rel=1e-9) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index df85decc676..d70a11dc1b9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26561,6 +26561,10 @@ export interface components { output_cost_per_pixel?: number | null; /** Output Cost Per Reasoning Token */ output_cost_per_reasoning_token?: number | null; + /** Output Cost Per Reasoning Token Flex */ + output_cost_per_reasoning_token_flex?: number | null; + /** Output Cost Per Reasoning Token Priority */ + output_cost_per_reasoning_token_priority?: number | null; /** Output Cost Per Second */ output_cost_per_second?: number | null; /** Output Cost Per Second 1080P */ @@ -35120,6 +35124,10 @@ export interface components { output_cost_per_pixel?: number | null; /** Output Cost Per Reasoning Token */ output_cost_per_reasoning_token?: number | null; + /** Output Cost Per Reasoning Token Flex */ + output_cost_per_reasoning_token_flex?: number | null; + /** Output Cost Per Reasoning Token Priority */ + output_cost_per_reasoning_token_priority?: number | null; /** Output Cost Per Second */ output_cost_per_second?: number | null; /** Output Cost Per Second 1080P */ From 2ba4e917666e84ca9b83d37ecbe807226e149020 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:38:08 -0700 Subject: [PATCH 04/16] feat(guardrails): add scan_only_tool_results to scope unified guardrails to tool results --- .../chat/guardrail_translation/handler.py | 35 ++++-- .../base_llm/guardrail_translation/utils.py | 53 +++++++- .../chat/guardrail_translation/handler.py | 33 +++-- .../proxy/guardrails/guardrail_registry.py | 12 +- litellm/types/guardrails.py | 10 ++ .../test_anthropic_guardrail_handler.py | 113 ++++++++++++++++++ .../test_openai_guardrail_handler.py | 68 +++++++++++ 7 files changed, 288 insertions(+), 36 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 3662389900b..535f4b7ae61 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -26,10 +26,10 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im ) from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, - openai_messages_without_system, - openai_messages_without_tool, + filtered_structured_messages, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -326,19 +326,25 @@ class AnthropicMessagesHandler(BaseTranslation): skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply) 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) chat_completion_compatible_request: Final = self._translate_to_openai(data) - structured_messages = cast( - list[AllMessageValues], - chat_completion_compatible_request.get("messages", []), + structured_messages: Final = list( + filtered_structured_messages( + cast( + list[AllMessageValues], + chat_completion_compatible_request.get("messages", []), + ), + scan_only_tool_results=scan_only_tool_results, + skip_system=skip_system, + skip_tool=skip_tool, + ) ) - if skip_system: - structured_messages = openai_messages_without_system(structured_messages) - if skip_tool: - structured_messages = openai_messages_without_tool(structured_messages) - tools_to_check: Final[list[ChatCompletionToolParam]] = chat_completion_compatible_request.get("tools", []) + tools_to_check: Final[list[ChatCompletionToolParam]] = ( + [] if scan_only_tool_results else chat_completion_compatible_request.get("tools", []) + ) # Step 1: Extract all text content and images extracted: Final = tuple( @@ -347,6 +353,7 @@ class AnthropicMessagesHandler(BaseTranslation): msg_idx=msg_idx, skip_system_message=skip_system, skip_tool_message=skip_tool, + scan_only_tool_results=scan_only_tool_results, ) for msg_idx, message in enumerate(messages) ) @@ -461,6 +468,7 @@ class AnthropicMessagesHandler(BaseTranslation): msg_idx: int, skip_system_message: bool = False, skip_tool_message: bool = False, + scan_only_tool_results: bool = False, ) -> ExtractedInput: """ Extract text content and images from a message. @@ -471,6 +479,8 @@ class AnthropicMessagesHandler(BaseTranslation): content: Final = message.get("content", None) if isinstance(content, str): + if scan_only_tool_results: + return EMPTY_EXTRACTED_INPUT return ExtractedInput(scanned=(ScannedText(content, MessageContentTarget(msg_idx)),), images=()) if not isinstance(content, list): return EMPTY_EXTRACTED_INPUT @@ -481,6 +491,7 @@ class AnthropicMessagesHandler(BaseTranslation): msg_idx=msg_idx, content_idx=content_idx, skip_tool_message=skip_tool_message, + scan_only_tool_results=scan_only_tool_results, ) for content_idx, content_item in enumerate(content) if isinstance(content_item, dict) @@ -497,12 +508,16 @@ class AnthropicMessagesHandler(BaseTranslation): msg_idx: int, content_idx: int, skip_tool_message: bool, + scan_only_tool_results: bool = False, ) -> ExtractedInput: if content_item.get("type") == "tool_result": if skip_tool_message: return EMPTY_EXTRACTED_INPUT return cls._extract_tool_result(content_item=content_item, msg_idx=msg_idx, content_idx=content_idx) + if scan_only_tool_results: + return EMPTY_EXTRACTED_INPUT + text_str: Final = content_item.get("text", None) return ExtractedInput( scanned=( diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 17cc0f118d6..e365913f2e1 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +from collections.abc import Sequence from typing import Any, Final from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage @@ -113,13 +114,53 @@ def effective_skip_tool_message_for_guardrail(guardrail_to_apply: Any) -> bool: return bool(getattr(litellm, "skip_tool_message_in_guardrail", False)) +def _message_role(message: AllMessageValues) -> str: + return str((message or {}).get("role") or "").lower() + + def openai_messages_without_system( - messages: list[AllMessageValues], -) -> list[AllMessageValues]: - return [m for m in messages if str((m or {}).get("role") or "").lower() != "system"] + messages: Sequence[AllMessageValues], +) -> tuple[AllMessageValues, ...]: + return tuple(m for m in messages if _message_role(m) != "system") def openai_messages_without_tool( - messages: list[AllMessageValues], -) -> list[AllMessageValues]: - return [m for m in messages if str((m or {}).get("role") or "").lower() != "tool"] + messages: Sequence[AllMessageValues], +) -> tuple[AllMessageValues, ...]: + return tuple(m for m in messages if _message_role(m) != "tool") + + +def openai_messages_only_tool( + messages: Sequence[AllMessageValues], +) -> tuple[AllMessageValues, ...]: + return tuple(m for m in messages if _message_role(m) == "tool") + + +def effective_scan_only_tool_results_for_guardrail(guardrail_to_apply: Any) -> bool: + return getattr(guardrail_to_apply, "scan_only_tool_results", None) is True + + +def role_out_of_guardrail_scope( + role: str, + *, + skip_system_message: bool, + skip_tool_message: bool, + scan_only_tool_results: bool = False, +) -> bool: + if skip_system_message and role == "system": + return True + if skip_tool_message and role == "tool": + return True + return scan_only_tool_results and role != "tool" + + +def filtered_structured_messages( + messages: Sequence[AllMessageValues], + *, + scan_only_tool_results: bool, + skip_system: bool, + skip_tool: bool, +) -> tuple[AllMessageValues, ...]: + scoped: Final = openai_messages_only_tool(messages) if scan_only_tool_results else tuple(messages) + without_system: Final = openai_messages_without_system(scoped) if skip_system else scoped + return openai_messages_without_tool(without_system) if skip_tool else without_system diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 3988326f2c2..67550890d2d 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -23,10 +23,11 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( StreamTransformSink, ) from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, - openai_messages_without_system, - openai_messages_without_tool, + filtered_structured_messages, + role_out_of_guardrail_scope, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -82,6 +83,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): skip_system: Final = effective_skip_system_message_for_guardrail(guardrail_to_apply) 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) texts_to_check: Final[list[str]] = [] images_to_check: Final[list[str]] = [] @@ -101,6 +103,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tool_call_task_mappings=tool_call_task_mappings, skip_system_message=skip_system, skip_tool_message=skip_tool, + scan_only_tool_results=scan_only_tool_results, ) # Step 2: Apply guardrail to all texts and tool calls in batch @@ -110,13 +113,16 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs["images"] = images_to_check if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check - structured_messages = self.get_structured_messages(data) + structured_messages: Final = self.get_structured_messages(data) if structured_messages: - if skip_system: - structured_messages = openai_messages_without_system(structured_messages) - if skip_tool: - structured_messages = openai_messages_without_tool(structured_messages) - inputs["structured_messages"] = structured_messages + inputs["structured_messages"] = list( + filtered_structured_messages( + structured_messages, + scan_only_tool_results=scan_only_tool_results, + skip_system=skip_system, + skip_tool=skip_tool, + ) + ) # Pass tools (function definitions) to the guardrail tools: Final = data.get("tools") if tools: @@ -194,16 +200,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tool_call_task_mappings: list[tuple[int, int]], skip_system_message: bool = False, skip_tool_message: bool = False, + scan_only_tool_results: bool = False, ) -> None: """ Extract text content, images, and tool calls from a message. Override this method to customize text/image/tool call extraction logic. """ - role: Final = str(message.get("role") or "").lower() - if skip_system_message and role == "system": - return - if skip_tool_message and role == "tool": + if role_out_of_guardrail_scope( + str(message.get("role") or "").lower(), + skip_system_message=skip_system_message, + skip_tool_message=skip_tool_message, + scan_only_tool_results=scan_only_tool_results, + ): return content: Final = message.get("content", None) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index f77588cf087..e9e61283c1a 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -487,16 +487,12 @@ class InMemoryGuardrailHandler: raise ValueError(f"Unsupported guardrail: {guardrail_type}") if custom_guardrail_callback is not None: - setattr( - custom_guardrail_callback, + for scoping_param in ( "skip_system_message_in_guardrail", - getattr(litellm_params, "skip_system_message_in_guardrail", None), - ) - setattr( - custom_guardrail_callback, "skip_tool_message_in_guardrail", - getattr(litellm_params, "skip_tool_message_in_guardrail", None), - ) + "scan_only_tool_results", + ): + setattr(custom_guardrail_callback, scoping_param, getattr(litellm_params, scoping_param, None)) configured_run_in_parallel: Final = getattr(litellm_params, "run_in_parallel", None) if configured_run_in_parallel is not None: custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index e7ad5cb801d..3eb8faf91dc 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -757,6 +757,16 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up ), ) + scan_only_tool_results: Optional[bool] = Field( + default=None, + description=( + "When True, unified guardrails only evaluate tool results, the untrusted data an " + "agent feeds back into the model, and skip system, user, and assistant content. " + "Intended for agent harnesses whose own prompt scaffolding is trusted but often " + "trips prompt-attack detectors." + ), + ) + # Lakera specific params category_thresholds: Optional[LakeraCategoryThresholds] = Field( default=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 dff3390af12..e90ae579d6d 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 @@ -760,3 +760,116 @@ class TestAnthropicMessagesToolResultScanning: assert "skip me POISON" not in guardrail.seen_texts assert messages[1]["content"][0]["content"] == "skip me POISON" assert messages[0]["content"] == "keep me [BLOCKED]" + + +class InputsRecordingGuardrail(MockMaskingGuardrail): + def __init__(self): + super().__init__(guardrail_name="scan-only-capture") + self.captured_inputs: Optional[GenericGuardrailAPIInputs] = None + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.captured_inputs = inputs + return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + + +class TestAnthropicMessagesScanOnlyToolResults: + def _guardrail(self): + guardrail = InputsRecordingGuardrail() + guardrail.scan_only_tool_results = True + return guardrail + + @pytest.mark.asyncio + async def test_scan_narrows_to_tool_results_and_write_back_stays_aligned(self): + handler = AnthropicMessagesHandler() + guardrail = self._guardrail() + data = { + "model": "claude-sonnet-4-5", + "system": "You are a trusted agent harness with POISON heuristics.", + "tools": [ + { + "name": "Bash", + "description": "run a command", + "input_schema": {"type": "object", "properties": {}}, + } + ], + "messages": [ + {"role": "user", "content": "scaffolding POISON prompt"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu1", "name": "Bash", "input": {"cmd": "curl"}}], + }, + { + "role": "user", + "content": [ + {"type": "text", "text": "sibling POISON text"}, + {"type": "tool_result", "tool_use_id": "tu1", "content": "fetched POISON page"}, + ], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.seen_texts == ["fetched POISON page"], ( + "only the tool_result payload may reach the guardrail" + ) + assert guardrail.captured_inputs is not None + assert guardrail.captured_inputs.get("tools") is None + assert [m["role"] for m in guardrail.captured_inputs["structured_messages"]] == ["tool"] + assert data["messages"][2]["content"][1]["content"] == "fetched [BLOCKED] page" + assert data["messages"][0]["content"] == "scaffolding POISON prompt", ( + "out-of-scope content must come back untouched, not masked or dropped" + ) + assert data["messages"][2]["content"][0]["text"] == "sibling POISON text" + + @pytest.mark.asyncio + async def test_guardrail_is_not_called_when_the_request_has_no_tool_results(self): + handler = AnthropicMessagesHandler() + guardrail = self._guardrail() + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "What is 2 plus 2?"}], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is None + assert guardrail.seen_texts == [] + + @pytest.mark.asyncio + async def test_images_are_scoped_the_same_way_as_texts(self): + handler = AnthropicMessagesHandler() + guardrail = self._guardrail() + data = { + "model": "claude-sonnet-4-5", + "messages": [ + { + "role": "user", + "content": [{"type": "image", "source": {"type": "base64", "data": "USER_IMG"}}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tu1", + "content": [ + {"type": "text", "text": "screenshot POISON"}, + {"type": "image", "source": {"type": "base64", "data": "TOOL_IMG"}}, + ], + } + ], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + assert guardrail.captured_inputs.get("images") == ["TOOL_IMG"] 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 7730b664c5e..c8a1b98aa82 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 @@ -1229,3 +1229,71 @@ class TestIncrementalScanRespectsSkipFlags: assert mock_api.call_count == 1 scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] assert scanned == ["It is sunny in Paris.", "And tomorrow?"] + + +class TestScanOnlyToolResults: + def _bedrock_guardrail(self): + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + guardrail = BedrockGuardrail( + guardrail_name="bedrock-scan-only-tool-results", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + ) + guardrail.scan_only_tool_results = True + return guardrail + + @pytest.mark.asyncio + async def test_only_tool_role_content_is_scanned(self): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + data = { + "messages": [ + {"role": "system", "content": "SYSTEM-PROMPT-not-scanned"}, + {"role": "user", "content": "USER-PROMPT-not-scanned"}, + { + "role": "assistant", + "content": "ASSISTANT-not-scanned", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "read_file", "arguments": '{"path": "report.html"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-scanned"}, + ] + } + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["TOOL-RESULT-scanned"] + + @pytest.mark.parametrize("flag_value", [None, "false", 0, object()]) + @pytest.mark.asyncio + async def test_scope_narrows_only_when_the_flag_is_actually_true(self, flag_value): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + guardrail.scan_only_tool_results = flag_value + data = { + "messages": [ + {"role": "user", "content": "USER-PROMPT"}, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}, + ] + } + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["USER-PROMPT", "TOOL-RESULT"], ( + "anything but an explicit True must leave the whole request in scope" + ) From d70e10982a46f728c6d5a431fd8692a85b3ebf23 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:21:58 -0700 Subject: [PATCH 05/16] fix(guardrails): keep tool-results-only scans off function definitions and merge scoped write-backs Gate the OpenAI handler's tools forwarding behind scan_only_tool_results, matching the Anthropic handler, so a tool-results-only scan can no longer evaluate or rewrite trusted function definitions. When a guardrail returns a replacement structured_messages list, substitute the returned messages back into the positions their scoped originals came from instead of installing the scoped list as the whole conversation, so out-of-scope messages (system prompt, prior turns) survive redaction on both the OpenAI and Anthropic paths. --- .../chat/guardrail_translation/handler.py | 32 +++++--- .../base_llm/guardrail_translation/utils.py | 58 ++++++++++--- .../chat/guardrail_translation/handler.py | 26 +++--- .../test_anthropic_guardrail_handler.py | 54 +++++++++++++ .../test_openai_guardrail_handler.py | 81 +++++++++++++++++++ 5 files changed, 216 insertions(+), 35 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 535f4b7ae61..c25fa624f7f 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -29,7 +29,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, - filtered_structured_messages, + merge_guardrailed_scoped_messages, + scoped_structured_message_indices, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -330,17 +331,17 @@ class AnthropicMessagesHandler(BaseTranslation): chat_completion_compatible_request: Final = self._translate_to_openai(data) - structured_messages: Final = list( - filtered_structured_messages( - cast( - list[AllMessageValues], - chat_completion_compatible_request.get("messages", []), - ), - scan_only_tool_results=scan_only_tool_results, - skip_system=skip_system, - skip_tool=skip_tool, - ) + full_structured_messages: Final = cast( + list[AllMessageValues], + chat_completion_compatible_request.get("messages", []), ) + scoped_message_indices: Final = scoped_structured_message_indices( + full_structured_messages, + scan_only_tool_results=scan_only_tool_results, + skip_system=skip_system, + skip_tool=skip_tool, + ) + structured_messages: Final = [full_structured_messages[index] for index in scoped_message_indices] tools_to_check: Final[list[ChatCompletionToolParam]] = ( [] if scan_only_tool_results else chat_completion_compatible_request.get("tools", []) @@ -402,7 +403,14 @@ class AnthropicMessagesHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): - self._write_back_structured_messages(data, guardrailed_structured_messages) + self._write_back_structured_messages( + data, + merge_guardrailed_scoped_messages( + full_messages=full_structured_messages, + scoped_indices=scoped_message_indices, + guardrailed_scoped=guardrailed_structured_messages, + ), + ) else: # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index e365913f2e1..fcd504fee2f 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from typing import Any, Final from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage @@ -130,12 +130,6 @@ def openai_messages_without_tool( return tuple(m for m in messages if _message_role(m) != "tool") -def openai_messages_only_tool( - messages: Sequence[AllMessageValues], -) -> tuple[AllMessageValues, ...]: - return tuple(m for m in messages if _message_role(m) == "tool") - - def effective_scan_only_tool_results_for_guardrail(guardrail_to_apply: Any) -> bool: return getattr(guardrail_to_apply, "scan_only_tool_results", None) is True @@ -154,13 +148,53 @@ def role_out_of_guardrail_scope( return scan_only_tool_results and role != "tool" -def filtered_structured_messages( +def scoped_structured_message_indices( messages: Sequence[AllMessageValues], *, scan_only_tool_results: bool, skip_system: bool, skip_tool: bool, -) -> tuple[AllMessageValues, ...]: - scoped: Final = openai_messages_only_tool(messages) if scan_only_tool_results else tuple(messages) - without_system: Final = openai_messages_without_system(scoped) if skip_system else scoped - return openai_messages_without_tool(without_system) if skip_tool else without_system +) -> tuple[int, ...]: + return tuple( + index + for index, message in enumerate(messages) + if not role_out_of_guardrail_scope( + _message_role(message), + skip_system_message=skip_system, + skip_tool_message=skip_tool, + scan_only_tool_results=scan_only_tool_results, + ) + ) + + +def merge_guardrailed_scoped_messages( + full_messages: Sequence[AllMessageValues], + scoped_indices: Sequence[int], + guardrailed_scoped: Sequence[AllMessageValues], +) -> list[AllMessageValues]: + """Substitute guardrail-returned messages back into the full conversation. + + Guardrails only ever see the scoped subset of messages, so a replacement + list they hand back describes that subset, not the whole request. Writing + it over ``data["messages"]`` wholesale would silently drop every + out-of-scope message (system prompt, prior turns). Instead, swap each + returned message into the position its scoped original came from; extra + returned messages land after the last scoped position, and scoped + originals without a counterpart are treated as removed by the guardrail. + When nothing was filtered out this degenerates to the returned list + itself, preserving wholesale-replacement behavior for unscoped guardrails. + """ + replacements: Final = dict(zip(scoped_indices, guardrailed_scoped)) + removed: Final = frozenset(scoped_indices[len(guardrailed_scoped) :]) + appended: Final = tuple(guardrailed_scoped[len(scoped_indices) :]) + last_scoped_index: Final = scoped_indices[-1] if scoped_indices else None + + def _merged() -> Iterator[AllMessageValues]: + for index, message in enumerate(full_messages): + if index in removed: + continue + yield replacements.get(index, message) + if index == last_scoped_index: + yield from appended + + return list(_merged()) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 67550890d2d..9d7fe6ce2a8 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -26,8 +26,9 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, - filtered_structured_messages, + merge_guardrailed_scoped_messages, role_out_of_guardrail_scope, + scoped_structured_message_indices, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -114,18 +115,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check structured_messages: Final = self.get_structured_messages(data) + scoped_message_indices: Final = scoped_structured_message_indices( + structured_messages or [], + scan_only_tool_results=scan_only_tool_results, + skip_system=skip_system, + skip_tool=skip_tool, + ) if structured_messages: - inputs["structured_messages"] = list( - filtered_structured_messages( - structured_messages, - scan_only_tool_results=scan_only_tool_results, - skip_system=skip_system, - skip_tool=skip_tool, - ) - ) + inputs["structured_messages"] = [structured_messages[index] for index in scoped_message_indices] # Pass tools (function definitions) to the guardrail tools: Final = data.get("tools") - if tools: + if tools and not scan_only_tool_results: inputs["tools"] = tools # Include model information if available model: Final = data.get("model") @@ -151,7 +151,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): - data["messages"] = guardrailed_structured_messages + data["messages"] = merge_guardrailed_scoped_messages( + full_messages=structured_messages or [], + scoped_indices=scoped_message_indices, + guardrailed_scoped=guardrailed_structured_messages, + ) else: # Step 3: Map guardrail responses back to original message structure if guardrailed_texts and texts_to_check: 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 e90ae579d6d..a016e1a2deb 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 @@ -5,6 +5,7 @@ Tests the handler's ability to process streaming output for Anthropic Messages A with guardrail transformations, specifically testing edge cases with empty choices. """ +import json import os import sys from typing import Any, Literal, Optional @@ -778,12 +779,65 @@ class InputsRecordingGuardrail(MockMaskingGuardrail): return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) +class StructuredMessagesRewritingGuardrail(CustomGuardrail): + """Returns a new structured_messages list with a canary redacted, like redaction guardrails do.""" + + def __init__(self): + super().__init__(guardrail_name="structured-rewrite") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + structured = inputs.get("structured_messages") or [] + inputs["structured_messages"] = [ + json.loads(json.dumps(message).replace("POISON", "[BLOCKED]")) for message in structured + ] + return inputs + + class TestAnthropicMessagesScanOnlyToolResults: def _guardrail(self): guardrail = InputsRecordingGuardrail() guardrail.scan_only_tool_results = True return guardrail + @pytest.mark.asyncio + async def test_structured_write_back_merges_into_the_full_conversation(self): + handler = AnthropicMessagesHandler() + guardrail = StructuredMessagesRewritingGuardrail() + guardrail.scan_only_tool_results = True + data = { + "model": "claude-sonnet-4-5", + "system": "You are a careful agent harness.", + "messages": [ + {"role": "user", "content": "fetch the page"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu1", "name": "Bash", "input": {"cmd": "curl"}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "tu1", "content": "fetched POISON page"}], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["system"] == "You are a careful agent harness." + assert [m["role"] for m in data["messages"]] == ["user", "assistant", "user"], ( + "a redacting guardrail must not strip out-of-scope turns from the request" + ) + serialized = json.dumps(data["messages"]) + assert "fetch the page" in serialized + assert "tool_use" in serialized + assert "fetched [BLOCKED] page" in serialized + assert "POISON" not in serialized + @pytest.mark.asyncio async def test_scan_narrows_to_tool_results_and_write_back_stays_aligned(self): handler = AnthropicMessagesHandler() 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 c8a1b98aa82..907da66e5bf 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 @@ -1231,6 +1231,28 @@ class TestIncrementalScanRespectsSkipFlags: assert scanned == ["It is sunny in Paris.", "And tomorrow?"] +class StructuredRedactionGuardrail(CustomGuardrail): + """Captures inputs and returns a new structured_messages list with a canary redacted.""" + + def __init__(self): + super().__init__(guardrail_name="structured-redaction") + self.captured_inputs: Optional[GenericGuardrailAPIInputs] = None + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.captured_inputs = inputs + structured = inputs.get("structured_messages") or [] + inputs["structured_messages"] = [ + {**m, "content": str(m.get("content", "")).replace("POISON", "[BLOCKED]")} for m in structured + ] + return inputs + + class TestScanOnlyToolResults: def _bedrock_guardrail(self): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail @@ -1297,3 +1319,62 @@ class TestScanOnlyToolResults: assert scanned == ["USER-PROMPT", "TOOL-RESULT"], ( "anything but an explicit True must leave the whole request in scope" ) + + @pytest.mark.parametrize("scan_only_tool_results", [True, False]) + @pytest.mark.asyncio + async def test_function_definitions_are_scoped_out_with_the_tool_results_flag(self, scan_only_tool_results): + handler = OpenAIChatCompletionsHandler() + guardrail = StructuredRedactionGuardrail() + guardrail.scan_only_tool_results = scan_only_tool_results + tools = [ + { + "type": "function", + "function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}}, + } + ] + data = { + "messages": [ + {"role": "user", "content": "read the report"}, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}, + ], + "tools": tools, + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert guardrail.captured_inputs is not None + expected_tools = None if scan_only_tool_results else tools + assert guardrail.captured_inputs.get("tools") == expected_tools, ( + "function definitions must stay out of a tool-results-only scan" + ) + + @pytest.mark.asyncio + async def test_structured_write_back_keeps_out_of_scope_messages(self): + handler = OpenAIChatCompletionsHandler() + guardrail = StructuredRedactionGuardrail() + guardrail.scan_only_tool_results = True + data = { + "messages": [ + {"role": "system", "content": "SYSTEM-PROMPT"}, + {"role": "user", "content": "fetch the page"}, + { + "role": "assistant", + "content": "fetching", + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "fetch", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"}, + {"role": "user", "content": "and then?"}, + ] + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [m["role"] for m in data["messages"]] == ["system", "user", "assistant", "tool", "user"], ( + "a redacting guardrail must not strip out-of-scope messages from the request" + ) + assert data["messages"][0]["content"] == "SYSTEM-PROMPT" + 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?" From 28a277e99e281052445e3568bcd1449a471447ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:56:50 -0700 Subject: [PATCH 06/16] refactor(guardrails): drop dead tool extraction and an Any annotation, ratchet lint budgets --- basedpyright-code-budget.json | 10 +++++----- .../chat/guardrail_translation/handler.py | 16 ---------------- .../llms/base_llm/guardrail_translation/utils.py | 2 +- ruff-strict-budget.json | 2 +- type-discipline-budget.json | 6 +++--- 5 files changed, 10 insertions(+), 26 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 27d96e415fd..8a5c78c1f6c 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 29204 }, "reportArgumentType": { - "limit": 2635 + "limit": 2634 }, "reportAssignmentType": { "limit": 329 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 9227 + "limit": 9226 }, "reportFunctionMemberAccess": { "limit": 7 @@ -105,7 +105,7 @@ "limit": 113 }, "reportUnknownMemberType": { - "limit": 40340 + "limit": 40339 }, "reportUnknownParameterType": { "limit": 20293 @@ -117,13 +117,13 @@ "limit": 122 }, "reportUnnecessaryComparison": { - "limit": 703 + "limit": 702 }, "reportUnnecessaryContains": { "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 865 + "limit": 864 }, "reportUntypedBaseClass": { "limit": 72 diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index c25fa624f7f..60424fb78b5 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -574,22 +574,6 @@ class AnthropicMessagesHandler(BaseTranslation): data: Final = source.get("data") return (data,) if data else () - def _extract_input_tools( - self, - tools: list[dict[str, Any]], - tools_to_check: list[ChatCompletionToolParam], - ) -> None: - """ - Extract tools from a message. - """ - ## CHECK FOR TOOLS - if tools is not None and isinstance(tools, list): - # TRANSFORM ANTHROPIC TOOLS TO OPENAI TOOLS - openai_tools: Final = self.adapter.translate_anthropic_tools_to_openai( - tools=cast(list[AllAnthropicToolsValues], tools) - ) - tools_to_check.extend(openai_tools) - async def _apply_guardrail_responses_to_input( self, messages: list[dict[str, Any]], diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index fcd504fee2f..432ac64b456 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -130,7 +130,7 @@ def openai_messages_without_tool( return tuple(m for m in messages if _message_role(m) != "tool") -def effective_scan_only_tool_results_for_guardrail(guardrail_to_apply: Any) -> bool: +def effective_scan_only_tool_results_for_guardrail(guardrail_to_apply: object) -> bool: return getattr(guardrail_to_apply, "scan_only_tool_results", None) is True diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 421b424757b..ea20ac97e07 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -42,7 +42,7 @@ "limit": 81 }, "B010": { - "limit": 194 + "limit": 192 }, "B018": { "limit": 2 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e26ce54ede7..37964c27657 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23343 + "limit": 23337 }, "LIT002": { "limit": 27213 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1093 + "limit": 1092 }, "LIT007": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16802 + "limit": 16796 }, "LIT011": { "limit": 5602 From 4b9872e7e890885d637bae32c62eb82998827cfe Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:53:44 -0700 Subject: [PATCH 07/16] fix(managed_files): return unified ids from unscoped file listing --- .../proxy/hooks/managed_files.py | 8 ++-- .../proxy/test_managed_files_hook.py | 39 +++++++++++++++++-- 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index ec47b6ac0e6..4c0b2b1d5cd 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -383,9 +383,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): } ) return [ - OpenAIFileObject.model_validate(file_object.file_object) - for file_object in file_ids - if file_object.file_object is not None + OpenAIFileObject.model_validate(row.file_object).model_copy( + update={"id": row.unified_file_id} + ) + for row in file_ids + if row.file_object is not None ] async def check_managed_file_id_access( diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 4a4aa7aa5ea..5dc98640449 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -142,8 +142,11 @@ async def test_get_user_created_file_ids_skips_rows_without_file_object(): managed_files = _make_managed_files_instance() managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( return_value=[ - MagicMock(file_object=_make_file_object().model_dump()), - MagicMock(file_object=None), + MagicMock( + file_object=_make_file_object().model_dump(), + unified_file_id="unified-id-1", + ), + MagicMock(file_object=None, unified_file_id="unified-id-2"), ] ) @@ -151,7 +154,37 @@ async def test_get_user_created_file_ids_skips_rows_without_file_object(): _make_user_api_key_dict(), ["file-output-abc"] ) - assert [file.id for file in files] == ["file-output-abc"] + assert [file.id for file in files] == ["unified-id-1"] + + +@pytest.mark.asyncio +async def test_get_user_created_file_ids_remaps_stored_raw_provider_id_to_unified_id(): + """ + Rows registered from batch outputs store the provider's file object, whose + id is the raw provider id (e.g. file-abc). Listing must return the row's + unified_file_id so callers get ids that work on the managed routes. + + Regression test for https://github.com/BerriAI/litellm/issues/35362. + """ + unified_id = "bGl0ZWxsbV9wcm94eTt1bmlmaWVkX2lkLGRlYWRiZWVm" + raw_provider_object = _make_file_object("file-raw-provider-123") + managed_files = _make_managed_files_instance() + managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[ + MagicMock( + file_object=raw_provider_object.model_dump(), + unified_file_id=unified_id, + ), + ] + ) + + files = await managed_files.get_user_created_file_ids( + _make_user_api_key_dict(), ["file-raw-provider-123"] + ) + + assert [file.id for file in files] == [unified_id] + assert files[0].filename == raw_provider_object.filename + assert files[0].purpose == raw_provider_object.purpose @pytest.mark.asyncio From c2998dea7510a3b656c06d54dbcdae769927b83d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:49:15 -0700 Subject: [PATCH 08/16] fix(guardrails): guard tools write-back under scan_only_tool_results and warn on role-filtered no-op scans --- .../chat/guardrail_translation/handler.py | 2 +- .../chat/guardrail_translation/handler.py | 2 +- .../guardrail_hooks/bedrock_guardrails.py | 8 +++ .../panw_prisma_airs/panw_prisma_airs.py | 12 +++++ .../test_anthropic_guardrail_handler.py | 34 ++++++++++++ .../test_openai_guardrail_handler.py | 54 +++++++++++++++++++ .../test_bedrock_guardrails.py | 37 +++++++++++++ .../guardrail_hooks/test_panw_prisma_airs.py | 28 ++++++++++ 8 files changed, 175 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 60424fb78b5..184e0f6a343 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -387,7 +387,7 @@ class AnthropicMessagesHandler(BaseTranslation): guardrailed_texts: Final = guardrailed_inputs.get("texts", []) guardrailed_tools: Final = guardrailed_inputs.get("tools") - if guardrailed_tools is not None: + if guardrailed_tools is not None and not scan_only_tool_results: # Convert tools back from OpenAI format to Anthropic format anthropic_config: Final = AnthropicConfig() anthropic_tools: Final[list[AllAnthropicToolsValues]] = [] diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 9d7fe6ce2a8..dc2a06d67fc 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -143,7 +143,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrailed_texts: Final = guardrailed_inputs.get("texts", []) guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", []) guardrailed_tools: Final = guardrailed_inputs.get("tools") - if guardrailed_tools is not None: + if guardrailed_tools is not None and not scan_only_tool_results: data["tools"] = guardrailed_tools guardrailed_structured_messages: Final = guardrailed_inputs.get("structured_messages") diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index f7a5c7559b1..8193069fd82 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -26,6 +26,9 @@ from litellm.caching import DualCache from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys +from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, +) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -523,6 +526,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): latest_user_index: Final = self._find_latest_message_index(structured_messages, target_role="user") if latest_user_index is None: + if effective_scan_only_tool_results_for_guardrail(self): + verbose_proxy_logger.warning( + "Bedrock Guardrail: experimental_use_latest_role_message_only scans only the latest " + "user message, so scan_only_tool_results leaves nothing to scan for this request" + ) verbose_proxy_logger.debug("Bedrock Guardrail: no user-role message in request, skipping INPUT scan") return ApplyGuardrailMessageSelection(None, None, True, skip_scan=True) 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 ae1478a9210..a96a0070eef 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 @@ -22,6 +22,9 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -1716,6 +1719,15 @@ class PanwPrismaAirsHandler(CustomGuardrail): # - latest-user extraction returned None (no user / count mismatch) if scannable_indices is None: scannable_indices = self._get_scannable_text_indices(texts, structured_messages) + if ( + scannable_indices is not None + and not scannable_indices + and effective_scan_only_tool_results_for_guardrail(self) + ): + verbose_proxy_logger.warning( + "PANW Prisma AIRS scans only user, system, and developer messages, " + "so scan_only_tool_results leaves nothing to scan for this request" + ) for i, text in enumerate(texts): if not text or not text.strip(): 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 a016e1a2deb..c7dedff0663 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 @@ -882,6 +882,40 @@ class TestAnthropicMessagesScanOnlyToolResults: ) assert data["messages"][2]["content"][0]["text"] == "sibling POISON text" + @pytest.mark.asyncio + async def test_guardrail_synthesized_tools_never_replace_scoped_out_request_tools(self): + handler = AnthropicMessagesHandler() + guardrail = ToolAppendingGuardrail(guardrail_name="tool-appending") + guardrail.scan_only_tool_results = True + original_tools = [ + { + "name": "get_weather", + "description": "Get the weather at a specific location", + "input_schema": {"type": "object", "properties": {"location": {"type": "string"}}}, + } + ] + data = { + "model": "claude-sonnet-4-5", + "tools": original_tools, + "messages": [ + {"role": "user", "content": "what's the weather?"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tu1", "name": "get_weather", "input": {}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "tu1", "content": "sunny"}], + }, + ], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["tools"] == original_tools, ( + "tools the guardrail synthesized without seeing the request's tools must not replace them" + ) + @pytest.mark.asyncio async def test_guardrail_is_not_called_when_the_request_has_no_tool_results(self): handler = AnthropicMessagesHandler() 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 907da66e5bf..269afef69cd 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 @@ -1253,6 +1253,31 @@ class StructuredRedactionGuardrail(CustomGuardrail): return inputs +class ToolSynthesizingGuardrail(CustomGuardrail): + """Appends its own function tool to whatever tools it was given, like a + retrieval/recovery guardrail that injects a tool the model can later call.""" + + def __init__(self): + super().__init__(guardrail_name="tool-synthesizing") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + tools = list(inputs.get("tools") or []) + tools.append( + { + "type": "function", + "function": {"name": "injected_retrieve", "parameters": {"type": "object", "properties": {}}}, + } + ) + inputs["tools"] = tools + return inputs + + class TestScanOnlyToolResults: def _bedrock_guardrail(self): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail @@ -1348,6 +1373,35 @@ class TestScanOnlyToolResults: "function definitions must stay out of a tool-results-only scan" ) + @pytest.mark.parametrize("scan_only_tool_results", [True, False]) + @pytest.mark.asyncio + async def test_guardrail_synthesized_tools_never_replace_scoped_out_request_tools(self, scan_only_tool_results): + handler = OpenAIChatCompletionsHandler() + guardrail = ToolSynthesizingGuardrail() + guardrail.scan_only_tool_results = scan_only_tool_results + original_tools = [ + { + "type": "function", + "function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}}, + } + ] + data = { + "messages": [ + {"role": "user", "content": "read the report"}, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}, + ], + "tools": original_tools, + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + if scan_only_tool_results: + assert data["tools"] == original_tools, ( + "tools the guardrail synthesized without seeing the request's tools must not replace them" + ) + else: + assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"] + @pytest.mark.asyncio async def test_structured_write_back_keeps_out_of_scope_messages(self): handler = OpenAIChatCompletionsHandler() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 65d6e33588f..76a695ce3fd 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -3670,3 +3670,40 @@ async def test_moderation_hook_honors_the_mcp_event_type(mode, call_type, should "the scan must be logged under the event it actually ran for, so guardrail logs, " "OTel spans, and Langfuse metadata do not misclassify MCP enforcement as an LLM call" ) + + +class TestScanOnlyToolResultsWithLatestRoleFilter: + @pytest.mark.asyncio + async def test_warns_and_skips_when_scoped_payload_has_no_user_message(self): + """scan_only_tool_results hands Bedrock a tool-role-only payload, but + experimental_use_latest_role_message_only scans only the latest user + message: the silent no-op must warn.""" + guardrail = BedrockGuardrail( + guardrail_name="bedrock-latest-role-scoped", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + default_on=True, + experimental_use_latest_role_message_only=True, + ) + guardrail.scan_only_tool_results = True + inputs = { + "texts": ["TOOL-RESULT"], + "structured_messages": [{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}], + } + + with ( + patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api, + patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.warning" + ) as mock_warning, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"litellm_call_id": "test-call-id"}, + input_type="request", + ) + + mock_api.assert_not_called() + assert result["texts"] == ["TOOL-RESULT"] + warning_text = " ".join(str(arg) for c in mock_warning.call_args_list for arg in c.args) + assert "scan_only_tool_results" in warning_text 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 431a7aa6f02..2f0fd51539d 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 @@ -1696,6 +1696,34 @@ class TestPanwAirsApplyGuardrail: request_data=request_data, guardrail_name=handler.guardrail_name ) + @pytest.mark.asyncio + async def test_apply_guardrail_warns_when_tool_results_scope_leaves_nothing_scannable(self, handler): + """scan_only_tool_results hands PANW a tool-role-only payload, but PANW's role + filter only scans user/system/developer rows: the silent no-op must warn.""" + handler.scan_only_tool_results = True + inputs: GenericGuardrailAPIInputs = { + "texts": ["TOOL-RESULT"], + "structured_messages": [{"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}], + } + request_data = {"litellm_call_id": "test-call-id", "model": "gpt-4"} + + with ( + patch.object(handler, "_call_panw_api", new_callable=AsyncMock) as mock_api, + patch( + "litellm.proxy.guardrails.guardrail_hooks.panw_prisma_airs.panw_prisma_airs.verbose_proxy_logger.warning" + ) as mock_warning, + ): + result = await handler.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + mock_api.assert_not_called() + assert result["texts"] == ["TOOL-RESULT"] + warning_text = " ".join(str(arg) for c in mock_warning.call_args_list for arg in c.args) + assert "scan_only_tool_results" in warning_text + @pytest.mark.asyncio async def test_apply_guardrail_block(self, handler): """Test block action raises HTTPException(400).""" From 0bae9708a729943a26b5e08313f08aabc5f3cab9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:57:50 -0700 Subject: [PATCH 09/16] fix(arize_phoenix): lowercase OTLP/gRPC auth metadata key (#34883) --- litellm/integrations/arize/arize_phoenix.py | 3 +- .../integrations/arize/test_arize_phoenix.py | 40 +++++++++++++++++-- 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index e13fc0184a4..5b52c59cae2 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -430,7 +430,8 @@ class ArizePhoenixLogger(OpenTelemetry): otlp_auth_headers = None if api_key is not None: - otlp_auth_headers = f"Authorization=Bearer {api_key}" + auth_header_key = "authorization" if protocol == "otlp_grpc" else "Authorization" + otlp_auth_headers = f"{auth_header_key}=Bearer {api_key}" elif "app.phoenix.arize.com" in endpoint: raise ValueError("PHOENIX_API_KEY must be set when using Phoenix Cloud (app.phoenix.arize.com).") diff --git a/tests/test_litellm/integrations/arize/test_arize_phoenix.py b/tests/test_litellm/integrations/arize/test_arize_phoenix.py index afd83f81ce0..9f79534242c 100644 --- a/tests/test_litellm/integrations/arize/test_arize_phoenix.py +++ b/tests/test_litellm/integrations/arize/test_arize_phoenix.py @@ -37,8 +37,8 @@ class TestArizePhoenixConfig(unittest.TestCase): # Call the function to get the configuration config = ArizePhoenixLogger.get_arize_phoenix_config() - # Verify the configuration - now uses standard Authorization Bearer format - self.assertEqual(config.otlp_auth_headers, "Authorization=Bearer test_api_key") + # gRPC metadata keys must be lowercase, so the auth header key is lowercased + self.assertEqual(config.otlp_auth_headers, "authorization=Bearer test_api_key") self.assertEqual(config.endpoint, "grpc://test.endpoint") self.assertEqual(config.protocol, "otlp_grpc") @@ -136,7 +136,7 @@ class TestArizePhoenixConfig(unittest.TestCase): "PHOENIX_COLLECTOR_ENDPOINT": "grpc://localhost:6006", "PHOENIX_API_KEY": "test_api_key", }, - "Authorization=Bearer test_api_key", + "authorization=Bearer test_api_key", "grpc://localhost:6006", "otlp_grpc", id="explicit grpc endpoint with grpc:// prefix", @@ -215,6 +215,40 @@ def test_get_arize_phoenix_config_expection_on_missing_api_key(monkeypatch, env_ ArizePhoenixLogger.get_arize_phoenix_config() +@pytest.mark.parametrize( + "collector_endpoint, expected_key", + [ + pytest.param("grpc://localhost:6006", "authorization", id="grpc prefix"), + pytest.param("http://localhost:4317", "authorization", id="grpc port 4317"), + pytest.param("http://localhost:6006", "Authorization", id="http"), + ], +) +def test_get_arize_phoenix_config_auth_header_key_casing( + monkeypatch, collector_endpoint, expected_key +): + """Regression for #34882: gRPC metadata keys must be lowercase. + + HTTP headers are case-insensitive, but the OTLP/gRPC exporter rejects an + uppercase ``Authorization`` metadata key, so span export silently fails. + """ + for key in [ + "PHOENIX_API_KEY", + "PHOENIX_COLLECTOR_ENDPOINT", + "PHOENIX_COLLECTOR_HTTP_ENDPOINT", + ]: + monkeypatch.delenv(key, raising=False) + + monkeypatch.setenv("PHOENIX_API_KEY", "test_api_key") + monkeypatch.setenv("PHOENIX_COLLECTOR_ENDPOINT", collector_endpoint) + + config = ArizePhoenixLogger.get_arize_phoenix_config() + + assert config.otlp_auth_headers == f"{expected_key}=Bearer test_api_key" + header_key = config.otlp_auth_headers.split("=", 1)[0] + if config.protocol == "otlp_grpc": + assert header_key == header_key.lower() + + # --------------------------------------------------------------------------- # Per-project routing via Resource (not span attributes) # --------------------------------------------------------------------------- From 7c621b31410dd116971db2c405d9cff55ed3b660 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 5 Aug 2026 21:03:36 -0700 Subject: [PATCH 10/16] fix(auto-router): accept every reminder marker pair a harness emits (#36029) * fix(auto-router): accept every reminder marker pair a harness emits reminder_markers held one (open, close) pair, so a harness that wraps injected context differently per agent type only got the slice of traffic using the configured envelope stripped. Every other agent type kept hitting the original bug: its reminder-only turn never stripped to empty, won "newest human ask", and the harness blob got classified in place of the real question, choosing the tier and therefore the spend. The field now takes a list of ReminderMarkerPair, following the KeywordTierRule pattern already in this file so each pair validates itself and errors point at reminder_markers.N.close rather than a bare index. Blocks from different pairs can nest, which the gap construction could not handle: resuming the kept text at an inner block's end walks back inside the enclosing block and leaks its remainder. Running the block ends through a maximum collapses nested and overlapping spans without a separate merge pass, and stays linear in block count, which a fold over a growing tuple of merged spans would not. A single pair's ends already increase, so the maximum is the identity and the default path is byte-identical: verified against the shipped function over 200k generated inputs, and every existing reminder test passes unchanged. The prior single-pair config shape is rejected loudly at startup and at /model/new rather than silently stripping nothing. * docs(auto-router): document reminder_markers in the complexity router README * chore(ui): regenerate dashboard API types for the reminder_markers shape --------- Co-authored-by: Abhimanyu Kapur <38531241+akapur99@users.noreply.github.com> --- .../complexity_router/README.md | 21 ++ .../complexity_router/__init__.py | 2 + .../complexity_router/complexity_router.py | 71 ++++-- .../complexity_router/config.py | 47 ++-- .../router_strategy/test_complexity_router.py | 222 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 27 ++- 6 files changed, 336 insertions(+), 54 deletions(-) diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index b1fdb0044be..259933dbb9e 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -171,6 +171,27 @@ If 2+ reasoning markers are detected in the user message, the request is automat Reasoning markers in the system prompt do **not** trigger the reasoning override. This prevents system prompts like "Think step by step before answering" from forcing all requests to the reasoning tier. +### Harness Reminder Blocks + +Agent harnesses inject their own context into the conversation as ordinary message text. That text is plumbing, not something a human asked for, so the router strips complete reminder blocks before classifying and picking a tier. A turn that is nothing but a reminder block strips to empty and is skipped, and the router falls back to the last real ask instead + +By default a block is anything between `` and ``. `reminder_markers` replaces that with your harness's own delimiters. Many harnesses use a different envelope per agent type, so list every pair you emit: + +```yaml +model_list: + - model_name: smart-router + litellm_params: + model: auto_router/complexity_router + complexity_router_config: + reminder_markers: + - open: "<<>>" + close: "<<>>" + - open: "[[SUBAGENT_CONTEXT_BEGIN]]" + close: "[[SUBAGENT_CONTEXT_END]]" +``` + +Setting `reminder_markers` replaces the built-in `` pair rather than adding to it, so list that pair too if your harness also emits it. Matching is case-insensitive. Blocks that nest or overlap across pairs are stripped whole. An unclosed delimiter is not a block and is left in place, which keeps prose that merely mentions a delimiter from being eaten + ### Code Detection Technical code keywords are detected case-insensitively and include: diff --git a/litellm/router_strategy/complexity_router/__init__.py b/litellm/router_strategy/complexity_router/__init__.py index 1830ff506e9..aa618cc807e 100644 --- a/litellm/router_strategy/complexity_router/__init__.py +++ b/litellm/router_strategy/complexity_router/__init__.py @@ -16,6 +16,7 @@ from litellm.router_strategy.complexity_router.config import ( DEFAULT_COMPLEXITY_CONFIG, ComplexityRouterConfig, ComplexityTier, + ReminderMarkerPair, ) __all__ = [ @@ -24,5 +25,6 @@ __all__ = [ "ComplexityRouter", "ComplexityRouterConfig", "ComplexityTier", + "ReminderMarkerPair", "classification_system_prompt", ] diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 7bbe01191e3..a69509fc37a 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -19,7 +19,7 @@ import asyncio import random import re from collections.abc import Iterator, Mapping, Sequence -from itertools import islice +from itertools import accumulate, islice from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast @@ -233,6 +233,7 @@ def _effective_turn_off_message_logging(request_kwargs: Mapping[str, Any] | None _REMINDER_OPEN: Final = "" _REMINDER_CLOSE: Final = "" +_DEFAULT_REMINDER_MARKERS: Final = ((_REMINDER_OPEN, _REMINDER_CLOSE),) _TRUNCATION_MARKER: Final = "..." @@ -253,10 +254,8 @@ def _message_text(content: object) -> str: return content if isinstance(content, str) else "" -def _reminder_block_spans( - lowered: str, open_marker: str = _REMINDER_OPEN, close_marker: str = _REMINDER_CLOSE -) -> Iterator[tuple[int, int]]: - """Span of each complete reminder block, left to right. +def _reminder_block_spans(lowered: str, open_marker: str, close_marker: str) -> Iterator[tuple[int, int]]: + """Span of each complete reminder block for one marker pair, left to right. Literal `str.find`, not a regex: the delimiters are fixed strings, and `.*?` retried its lazy quantifier from every opening tag, so repeated unclosed tags were quadratic @@ -272,17 +271,36 @@ def _reminder_block_spans( yield start, cursor -def _strip_reminder_blocks(text: str, open_marker: str = _REMINDER_OPEN, close_marker: str = _REMINDER_CLOSE) -> str: - """Remove every complete reminder block from text, keeping everything written around them.""" - spans: Final = tuple(_reminder_block_spans(text.lower(), open_marker, close_marker)) +def _strip_reminder_blocks(text: str, marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS) -> str: + """Remove every complete reminder block from text, keeping everything written around them. + + Blocks from different pairs can nest or overlap, which the gap construction below would + otherwise mishandle: an inner block's end would resume the kept text partway through the outer + block, leaking the rest of that block into the classified ask. Running the block ends through a + maximum resumes each gap past the furthest block seen so far, which collapses nested and + overlapping spans without a separate merge pass. A single pair's ends already increase, so the + maximum is the identity there and the default path is byte-identical to a plain scan. + + Deliberately linear in both the text and the block count. This runs pre-routing on input any + keyholder controls, and both a regex scan and a fold that rebuilds a growing tuple of merged + spans go quadratic on inputs that are cheap to send. + """ + lowered: Final = text.lower() + spans: Final = tuple( + sorted( + span + for open_marker, close_marker in marker_pairs + for span in _reminder_block_spans(lowered, open_marker, close_marker) + ) + ) if not spans: return text.strip() - keep_from: Final = (0, *(end for _, end in spans)) + keep_from: Final = (0, *accumulate((end for _, end in spans), max)) keep_to: Final = (*(start for start, _ in spans), len(text)) return " ".join(kept for a, b in zip(keep_from, keep_to) if (kept := text[a:b].strip())) -def _human_text(content: object, open_marker: str = _REMINDER_OPEN, close_marker: str = _REMINDER_CLOSE) -> str: +def _human_text(content: object, marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS) -> str: """Message content as the text a human wrote, with complete reminder blocks removed. Harnesses inject reminders as ordinary text alongside the live ask, so the block is stripped and @@ -291,18 +309,18 @@ def _human_text(content: object, open_marker: str = _REMINDER_OPEN, close_marker one, and this same string drives escalation keywords and keyword_tier_rules, which choose the model and therefore the spend. An unclosed tag is not a block and is left intact. """ - return _strip_reminder_blocks(_message_text(content), open_marker, close_marker) + return _strip_reminder_blocks(_message_text(content), marker_pairs) def _iter_human_asks_newest_first( - messages: Sequence[Mapping[str, object]], markers: tuple[str, str] = (_REMINDER_OPEN, _REMINDER_CLOSE) + messages: Sequence[Mapping[str, object]], + marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, ) -> Iterator[str]: """Yield user-turn texts that carry a real human ask, newest first, with harness noise removed.""" - open_marker, close_marker = markers return ( text for msg in reversed(messages) - if msg.get("role") == "user" and (text := _human_text(msg.get("content"), open_marker, close_marker)) + if msg.get("role") == "user" and (text := _human_text(msg.get("content"), marker_pairs)) ) @@ -341,7 +359,8 @@ def _conversation_is_continuing(messages: Sequence[Mapping[str, object]] | None) def _newest_turn_ask( - messages: Sequence[Mapping[str, object]], markers: tuple[str, str] = (_REMINDER_OPEN, _REMINDER_CLOSE) + messages: Sequence[Mapping[str, object]], + marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, ) -> str | None: """The human ask on the newest user turn, or None when that turn carries only plumbing. @@ -352,12 +371,12 @@ def _newest_turn_ask( newest_user_turn: Final = next((msg for msg in reversed(messages) if msg.get("role") == "user"), None) if newest_user_turn is None: return None - return _human_text(newest_user_turn.get("content"), *markers) or None + return _human_text(newest_user_turn.get("content"), marker_pairs) or None def _extract_current_ask_and_system_prompt( messages: Sequence[Mapping[str, object]], - markers: tuple[str, str] = (_REMINDER_OPEN, _REMINDER_CLOSE), + marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, ) -> tuple[str | None, str | None]: """The last real human ask and the last system prompt; either is None if absent. @@ -365,7 +384,7 @@ def _extract_current_ask_and_system_prompt( the caller routes to its default model. That is the correct answer rather than a gap to fill: filling it would hand tier selection to harness-injected text. """ - current_ask: Final = next(_iter_human_asks_newest_first(messages, markers), None) + current_ask: Final = next(_iter_human_asks_newest_first(messages, marker_pairs), None) system_prompt: Final = next( ( text @@ -385,7 +404,7 @@ def _truncate(text: str, limit: int) -> str: def _iter_context_turns_newest_first( messages: Sequence[Mapping[str, object]], include_assistant: bool, - markers: tuple[str, str] = (_REMINDER_OPEN, _REMINDER_CLOSE), + marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, ) -> Iterator[tuple[str, str]]: """Yield (role, text) for turns eligible as classifier context, newest first. @@ -401,7 +420,7 @@ def _iter_context_turns_newest_first( for msg in reversed(messages) if isinstance(role := msg.get("role"), str) and role in roles - and (text := _human_text(msg.get("content"), *markers)) + and (text := _human_text(msg.get("content"), marker_pairs)) ) @@ -411,7 +430,7 @@ def _extract_prior_turns( window_size: int, per_turn_chars: int, include_assistant: bool, - markers: tuple[str, str] = (_REMINDER_OPEN, _REMINDER_CLOSE), + marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, ) -> tuple[tuple[str, str], ...]: """Up to window_size turns other than current_ask, oldest first, as (role, text). @@ -431,7 +450,7 @@ def _extract_prior_turns( prior: Final = islice( ( turn - for turn in _iter_context_turns_newest_first(messages, include_assistant, markers) + for turn in _iter_context_turns_newest_first(messages, include_assistant, marker_pairs) if turn[1] != current_ask ), window_size, @@ -556,7 +575,11 @@ class ComplexityRouter(CustomLogger): if self.config.escalation_keywords is not None else DEFAULT_ESCALATION_KEYWORDS ) - self._reminder_markers: tuple[str, str] = self.config.reminder_markers or (_REMINDER_OPEN, _REMINDER_CLOSE) + self._reminder_markers: tuple[tuple[str, str], ...] = ( + tuple((pair.open, pair.close) for pair in self.config.reminder_markers) + if self.config.reminder_markers + else _DEFAULT_REMINDER_MARKERS + ) # Lazily built on first semantic request and cached for reuse (route # embeddings are static, only the prompt is embedded per request). The lock @@ -993,7 +1016,7 @@ class ComplexityRouter(CustomLogger): window_size=self.config.classifier_context_window_size, per_turn_chars=self.config.classifier_context_per_turn_chars, include_assistant=include_assistant, - markers=self._reminder_markers, + marker_pairs=self._reminder_markers, ) if context_enabled else () diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index f9d3bd9ae67..69609a973b0 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -59,6 +59,30 @@ class KeywordTierRule(BaseModel): return self +class ReminderMarkerPair(BaseModel): + """One open/close delimiter pair a harness wraps injected context in. + + Normalizing here rather than at the scan is what makes matching case-insensitive: markers reach + the scan already lowered, so it lowercases only the haystack and never the needles. Stripping + keeps YAML indentation whitespace from becoming part of the delimiter. + """ + + open: str = Field(description="Opening delimiter, e.g. ''") + close: str = Field(description="Closing delimiter, e.g. ''") + + @model_validator(mode="after") + def _normalize(self) -> "ReminderMarkerPair": + open_marker: Final = self.open.strip().lower() + close_marker: Final = self.close.strip().lower() + if not open_marker or not close_marker: + raise ValueError("reminder_markers entries must not be blank") + if open_marker == close_marker: + raise ValueError("reminder_markers open and close must be different strings") + self.open = open_marker + self.close = close_marker + return self + + # ─── Default Keyword Lists ─── # Note: Keywords should be full words/phrases to avoid substring false positives. # The matching logic uses word boundary detection for single-word keywords. @@ -498,12 +522,15 @@ class ComplexityRouterConfig(BaseModel): description="RoutingPlugin instances that narrow the classified tier's candidate models before selection", ) - reminder_markers: tuple[str, str] | None = Field( + reminder_markers: tuple[ReminderMarkerPair, ...] | None = Field( default=None, + min_length=1, description=( - "Override the (open, close) marker pair used to recognize and strip harness-injected " - "reminder blocks before classification. Defaults to Claude Code's convention, " - "('', ''), when unset. Matching is case-insensitive." + "Override the delimiter pairs used to recognize and strip harness-injected reminder " + "blocks before classification. A harness that wraps injected context differently per " + "agent type (main, subagent, cron) lists every pair it emits. Replaces, rather than " + "adds to, the built-in default of ('', ''), so a " + "harness that also emits that pair lists it too. Matching is case-insensitive." ), ) @@ -601,18 +628,6 @@ class ComplexityRouterConfig(BaseModel): ) return self - @model_validator(mode="after") - def _normalize_reminder_markers(self) -> "ComplexityRouterConfig": - if self.reminder_markers is None: - return self - open_marker, close_marker = (marker.strip().lower() for marker in self.reminder_markers) - if not open_marker or not close_marker: - raise ValueError("reminder_markers entries must not be blank") - if open_marker == close_marker: - raise ValueError("reminder_markers open and close must be different strings") - self.reminder_markers = (open_marker, close_marker) - return self - def tier_label(self, tier: ComplexityTier) -> str: """Operator-facing display name for a tier, falling back to its canonical name.""" return self.tier_labels.get(tier, "").strip() or tier.value diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index b3f1e929741..8e9e32f5898 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -112,6 +112,52 @@ class TestComplexityRouterInit: assert router.config.tiers["SIMPLE"] == "gpt-4o-mini" assert router.config.tiers["REASONING"] == "o1-preview" + def test_configured_marker_pairs_reach_the_ask_extraction(self, mock_router_instance, basic_config): + """Marker pairs configured in YAML must actually reach the code that strips them. + + The config field, the validator and the scan were each covered on their own, but nothing + exercised config.reminder_markers -> self._reminder_markers, so the router could have parsed + a valid config and still classified on unstripped text. Asserting through the extraction the + router feeds its classifier is what makes that wiring a regression rather than a silent gap. + """ + from litellm.router_strategy.complexity_router.complexity_router import ( + _extract_current_ask_and_system_prompt, + ) + + ask = "Derive the amortized complexity of a splay tree access" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + **basic_config, + "reminder_markers": [ + {"open": "<<>>", "close": "<<>>"}, + {"open": "[[SUBAGENT_BEGIN]]", "close": "[[SUBAGENT_END]]"}, + ], + }, + ) + + assert router._reminder_markers == ( + ("<<>>", "<<>>"), + ("[[subagent_begin]]", "[[subagent_end]]"), + ) + messages = [ + {"role": "user", "content": ask}, + {"role": "assistant", "content": "Working on it."}, + {"role": "user", "content": "[[SUBAGENT_BEGIN]]Budget: 42 tokens remaining.[[SUBAGENT_END]]"}, + ] + assert _extract_current_ask_and_system_prompt(messages, router._reminder_markers)[0] == ask + + def test_unconfigured_marker_pairs_fall_back_to_the_builtin_default(self, mock_router_instance, basic_config): + """A config that never mentions reminder_markers keeps stripping .""" + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=basic_config, + ) + + assert router._reminder_markers == (("", ""),) + def test_init_without_config(self, mock_router_instance): """Test initialization without configuration uses defaults.""" router = ComplexityRouter( @@ -2991,17 +3037,68 @@ class TestSemanticConfigValidation: def test_reminder_markers_are_normalized(self): """Markers are stripped and lowercased, matching how the built-in constants are compared.""" config = ComplexityRouterConfig( - reminder_markers=(" <<>> ", "<<>>"), + reminder_markers=[{"open": " <<>> ", "close": "<<>>"}], ) - assert config.reminder_markers == ("<<>>", "<<>>") + assert config.reminder_markers is not None + assert (config.reminder_markers[0].open, config.reminder_markers[0].close) == ( + "<<>>", + "<<>>", + ) + + def test_reminder_markers_keep_every_configured_pair_in_order(self): + """Every pair a harness emits survives validation, not just the first.""" + config = ComplexityRouterConfig( + reminder_markers=[ + {"open": "<<>>", "close": "<<>>"}, + {"open": "[[SUBAGENT_BEGIN]]", "close": "[[SUBAGENT_END]]"}, + {"open": "%%CRON_BEGIN%%", "close": "%%CRON_END%%"}, + ], + ) + assert config.reminder_markers is not None + assert [(pair.open, pair.close) for pair in config.reminder_markers] == [ + ("<<>>", "<<>>"), + ("[[subagent_begin]]", "[[subagent_end]]"), + ("%%cron_begin%%", "%%cron_end%%"), + ] def test_reminder_markers_reject_blank_entry(self): with pytest.raises(ValidationError, match="must not be blank"): - ComplexityRouterConfig(reminder_markers=("", "<<>>")) + ComplexityRouterConfig(reminder_markers=[{"open": "", "close": "<<>>"}]) def test_reminder_markers_reject_identical_open_and_close(self): with pytest.raises(ValidationError, match="must be different"): - ComplexityRouterConfig(reminder_markers=("<<>>", "<<>>")) + ComplexityRouterConfig(reminder_markers=[{"open": "<<>>", "close": "<<>>"}]) + + def test_reminder_markers_reject_a_bad_pair_anywhere_in_the_list(self): + """Validation runs per pair, so a broken entry after a good one is still caught.""" + with pytest.raises(ValidationError, match="must be different"): + ComplexityRouterConfig( + reminder_markers=[ + {"open": "<<>>", "close": "<<>>"}, + {"open": "<<>>", "close": "<<>>"}, + ], + ) + + def test_reminder_markers_reject_empty_list(self): + """An explicitly empty list is ambiguous, so it fails loudly instead of silently defaulting. + + Left to fall through, an empty list resolves to the built-in pair, which + reads as "strip nothing" in the config and does the opposite. Matching on the length error + keeps this from passing for some unrelated reason if the field type changes. + """ + with pytest.raises(ValidationError, match="at least 1 item"): + ComplexityRouterConfig(reminder_markers=[]) + + def test_reminder_markers_reject_the_old_flat_pair_form(self): + """The pre-list shape is rejected loudly rather than silently routing on unstripped text. + + reminder_markers took a bare (open, close) string pair before it took a list of pairs. A + config still using that shape must fail validation at startup and at /model/new write time, + because the alternative -- accepting it and stripping nothing -- hands tier selection, and + therefore spend, to harness-injected text without any signal that it happened. + """ + with pytest.raises(ValidationError, match="valid dictionary or instance of ReminderMarkerPair"): + ComplexityRouterConfig(reminder_markers=("", "")) class _StubEncoder: @@ -4306,7 +4403,6 @@ class TestRoutingDecisionContents: # The score is still recorded, but the cause is what says it did not decide. assert decision["score"] < decision["tier_boundaries"]["complex_reasoning"] - @pytest.mark.asyncio async def test_an_unrenamed_router_writes_no_tier_label(self, complexity_router): """Renaming is opt-in, so a deployment that never renamed must gain no new key. @@ -4919,12 +5015,73 @@ class TestContextAwareClassifier: """ from litellm.router_strategy.complexity_router.complexity_router import _extract_current_ask_and_system_prompt - markers = ("<<>>", "<<>>") - follow_up_reminder = f"{markers[0]}Budget: 42 tokens remaining. Do not mention this.{markers[1]}" + pair = ("<<>>", "<<>>") + follow_up_reminder = f"{pair[0]}Budget: 42 tokens remaining. Do not mention this.{pair[1]}" messages = [_ASKED, _ANSWERED, {"role": "user", "content": follow_up_reminder}] assert _extract_current_ask_and_system_prompt(messages)[0] == follow_up_reminder - assert _extract_current_ask_and_system_prompt(messages, markers)[0] == _ASK + assert _extract_current_ask_and_system_prompt(messages, (pair,))[0] == _ASK + + def test_every_configured_marker_pair_is_stripped_not_just_the_first(self): + """One deployment serves a harness whose agent types each use a different envelope. + + Main agent, subagent and cron wrap injected context in different open/close pairs, and they + all route through the same auto-router. When only one pair could be configured, the other + agent types kept hitting the original bug: their reminder-only turn never stripped to empty, + won "newest human ask", and the harness blob got classified in place of the real question. + Each pair in turn must be skipped, so this fails if only the first configured pair is used. + """ + from litellm.router_strategy.complexity_router.complexity_router import _extract_current_ask_and_system_prompt + + pairs = ( + ("<<>>", "<<>>"), + ("[[subagent_begin]]", "[[subagent_end]]"), + ("%%cron_begin%%", "%%cron_end%%"), + ) + for open_marker, close_marker in pairs: + reminder_only_turn = f"{open_marker}Budget: 42 tokens remaining.{close_marker}" + messages = [_ASKED, _ANSWERED, {"role": "user", "content": reminder_only_turn}] + + assert _extract_current_ask_and_system_prompt(messages, pairs)[0] == _ASK, open_marker + + def test_a_block_nested_inside_another_pairs_block_does_not_leak(self): + """Nested blocks from two pairs must strip whole, not resume inside the outer block. + + Spans are collected per pair and can nest. Resuming the kept text at each block's own end + walks backwards into the enclosing block, so the outer block's remainder (and its dangling + close marker) survive into the classified ask. That is harness text choosing the tier, and + therefore the spend. Overlapping and disjoint spans strip correctly either way, so this + nested case is what pins the behavior. + """ + from litellm.router_strategy.complexity_router.complexity_router import _strip_reminder_blocks + + pairs = (("<<>>", "<<>>"), ("[[subagent_begin]]", "[[subagent_end]]")) + nested = "<<>>budget[[subagent_begin]]inner[[subagent_end]]do not mention<<>>" + + assert _strip_reminder_blocks(f"{nested} what is a splay tree?", pairs) == "what is a splay tree?" + + def test_overlapping_blocks_from_two_pairs_strip_whole(self): + """Interleaved (not nested) blocks still strip everything they jointly cover.""" + from litellm.router_strategy.complexity_router.complexity_router import _strip_reminder_blocks + + pairs = (("<<>>", "<<>>"), ("[[subagent_begin]]", "[[subagent_end]]")) + overlapping = "<<>>a[[subagent_begin]]b<<>>c[[subagent_end]]" + + assert _strip_reminder_blocks(f"{overlapping} what is a splay tree?", pairs) == "what is a splay tree?" + + def test_an_unclosed_marker_in_one_pair_does_not_suppress_another_pairs_blocks(self): + """Each pair scans independently, so one pair's dangling opener is not a global stop. + + An unclosed tag ends that pair's scan by design and is left intact as prose. It must not + also swallow a different pair's complete block, which would put harness text back in front + of the classifier. + """ + from litellm.router_strategy.complexity_router.complexity_router import _strip_reminder_blocks + + pairs = (("<<>>", "<<>>"), ("[[subagent_begin]]", "[[subagent_end]]")) + text = "<<>> why is [[subagent_begin]]noise[[subagent_end]] my tag stripped?" + + assert _strip_reminder_blocks(text, pairs) == "<<>> why is my tag stripped?" @pytest.mark.parametrize( "messages,current_ask,window,per_turn_chars,include_assistant,expected", @@ -5084,6 +5241,28 @@ class TestContextAwareClassifier: assert _extract_prior_turns(messages, current_ask, window, per_turn_chars, include_assistant) == expected + def test_prior_turn_context_strips_every_configured_pair(self): + """The classifier's context window is stripped with the same pairs as the ask. + + Prior turns are quoted verbatim into the LLM classifier payload, so a pair that is honored + when picking the ask but ignored when building context puts the harness blob back in front + of the classifier through the other door. This covers the _extract_prior_turns call the ask + extraction tests never reach. + """ + from litellm.router_strategy.complexity_router.complexity_router import _extract_prior_turns + + pairs = (("<<>>", "<<>>"), ("[[subagent_begin]]", "[[subagent_end]]")) + messages = [ + {"role": "user", "content": "[[subagent_begin]]budget blob[[subagent_end]]what about b-trees?"}, + {"role": "user", "content": "<<>>other blob<<>>and heaps?"}, + {"role": "user", "content": "current ask"}, + ] + + assert _extract_prior_turns(messages, "current ask", 5, 200, False, pairs) == ( + ("user", "what about b-trees?"), + ("user", "and heaps?"), + ) + def test_reminder_scan_is_linear_on_adversarial_input(self): """Unclosed reminder tags must not make stripping superlinear. @@ -5105,6 +5284,29 @@ class TestContextAwareClassifier: assert elapsed < 1.0, f"stripping {len(adversarial)} chars took {elapsed:.2f}s; scan is not linear" assert result == adversarial + def test_reminder_scan_stays_linear_in_block_count_across_pairs(self): + """Many *complete* blocks across several pairs must not go quadratic either. + + Collapsing nested and overlapping spans is required for correctness once more than one pair + is configured, and the obvious way to write it -- folding merged spans into a growing tuple + -- is quadratic in block count. Unlike the unclosed-tag case above, these blocks all close, + so they actually produce spans. This input is a few hundred KB, which any keyholder can send + pre-routing, and it fails loudly if the collapse is ever rewritten as a fold. + """ + import time + + from litellm.router_strategy.complexity_router.complexity_router import _strip_reminder_blocks + + pairs = (("", ""), ("", "")) + adversarial = "xy" * 25_000 + + start = time.perf_counter() + result = _strip_reminder_blocks(f"{adversarial} what is a splay tree?", pairs) + elapsed = time.perf_counter() - start + + assert elapsed < 1.0, f"stripping {50_000} blocks took {elapsed:.2f}s; collapse is not linear" + assert result == "what is a splay tree?" + @pytest.mark.asyncio async def test_llm_classifier_includes_prior_turns_context(self, llm_complexity_router, mock_router_instance): """Test that the LLM classifier receives prior-turn context in the user message.""" @@ -5761,7 +5963,9 @@ class TestCustomClassifierSystemPrompt: @pytest.mark.asyncio async def test_custom_prompt_is_sent_verbatim_as_the_system_role(self, mock_router_instance, llm_classifier_config): - custom = "Classify the data sensitivity: SIMPLE=public, MEDIUM=internal, COMPLEX=confidential, REASONING=regulated." + custom = ( + "Classify the data sensitivity: SIMPLE=public, MEDIUM=internal, COMPLEX=confidential, REASONING=regulated." + ) router = ComplexityRouter( model_name="test-complexity-router", litellm_router_instance=mock_router_instance, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 65434407f74..959c9921896 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -31560,6 +31560,26 @@ export interface components { /** Review Notes */ review_notes?: string | null; }; + /** + * ReminderMarkerPair + * @description One open/close delimiter pair a harness wraps injected context in. + * + * Normalizing here rather than at the scan is what makes matching case-insensitive: markers reach + * the scan already lowered, so it lowercases only the haystack and never the needles. Stripping + * keeps YAML indentation whitespace from becoming part of the delimiter. + */ + ReminderMarkerPair: { + /** + * Close + * @description Closing delimiter, e.g. '' + */ + close: string; + /** + * Open + * @description Opening delimiter, e.g. '' + */ + open: string; + }; /** * RequestComplexityRouterConfig * @description The part of a complexity-router config a request can carry. @@ -31672,12 +31692,9 @@ export interface components { reasoning_keywords?: string[] | null; /** * Reminder Markers - * @description Override the (open, close) marker pair used to recognize and strip harness-injected reminder blocks before classification. Defaults to Claude Code's convention, ('', ''), when unset. Matching is case-insensitive. + * @description Override the delimiter pairs used to recognize and strip harness-injected reminder blocks before classification. A harness that wraps injected context differently per agent type (main, subagent, cron) lists every pair it emits. Replaces, rather than adds to, the built-in default of ('', ''), so a harness that also emits that pair lists it too. Matching is case-insensitive. */ - reminder_markers?: [ - string, - string - ] | null; + reminder_markers?: components["schemas"]["ReminderMarkerPair"][] | null; /** * Return Raw Model Name * @description Return the resolved raw model name in the response model field instead of the client-requested complexity-router alias From 3d275d97feacc8a0e2a0d35bbd0f108a93e9971d Mon Sep 17 00:00:00 2001 From: Michael Cusack Date: Wed, 5 Aug 2026 21:49:31 -0700 Subject: [PATCH 11/16] fix(router): return model and Bedrock batch fields in deployment credentials get_deployment_credentials_with_provider dropped s3_region_name, s3_encryption_key_id, and aws_batch_role_arn because CredentialLiteLLMParams never declared them, and it never returned the deployment's model, so proxy batch creation against Bedrock failed with "LiteLLM doesn't support custom_llm_provider=bedrock for 'create_batch'" or "AWS IAM role ARN is required" (#25104) Provider-only file and batch calls keep their no-model contract: get_team_provider_credentials strips the model key so a provider-scoped request is not pinned to an arbitrary matching deployment --- .../openai_files_endpoints/common_utils.py | 2 +- litellm/router.py | 4 ++- litellm/types/router.py | 8 ++--- tests/test_litellm/test_router.py | 36 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 +++++ 5 files changed, 51 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 080b8b80ae4..0f6494051cf 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -373,7 +373,7 @@ def get_team_provider_credentials( def _provider_credentials(model_id: str) -> dict | None: credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id) if credentials is not None and credentials.get("custom_llm_provider") == custom_llm_provider: - return credentials + return {key: value for key, value in credentials.items() if key != "model"} return None # 1. Prefer the team's own BYOK deployment, matched by model_info.team_id. diff --git a/litellm/router.py b/litellm/router.py index c4a16521fa7..1edb80da7ce 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8738,7 +8738,7 @@ class Router: Example: credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm") - # Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", ...} + # Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", "model": "gpt-4o", ...} """ # Try to get deployment by model_id first deployment = self.get_deployment(model_id=model_id) @@ -8797,6 +8797,8 @@ class Router: # Remove the credential name since we've resolved it credentials.pop("litellm_credential_name", None) + credentials["model"] = deployment.litellm_params.model + # Add custom_llm_provider if deployment.litellm_params.custom_llm_provider: credentials["custom_llm_provider"] = deployment.litellm_params.custom_llm_provider diff --git a/litellm/types/router.py b/litellm/types/router.py index 83757daa4dd..8b4b547bdcc 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -200,6 +200,9 @@ class CredentialLiteLLMParams(BaseModel): aws_bedrock_runtime_endpoint: str | None = None aws_bedrock_project_id: str | None = None s3_bucket_name: str | None = None + s3_region_name: str | None = None + s3_encryption_key_id: str | None = None + aws_batch_role_arn: str | None = None ## IBM WATSONX ## watsonx_region_name: str | None = None @@ -272,11 +275,6 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): quality_router_config: dict | None = None quality_router_default_model: str | None = None - # Batch/File API Params - s3_bucket_name: str | None = None - s3_encryption_key_id: str | None = None - gcs_bucket_name: str | None = None - # Vector Store Params vector_store_id: str | None = None milvus_text_field: str | None = None diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 4dec574b9f3..b910cc3c5fc 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -4024,6 +4024,42 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): litellm.credential_list = [] +def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): + """ + Test that get_deployment_credentials_with_provider returns the deployment's + model and the Bedrock batch/S3 fields (s3_region_name, s3_encryption_key_id, + aws_batch_role_arn) instead of silently dropping them (#25104). + """ + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-batch-model", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-batch-bucket", + "s3_region_name": "us-east-1", + "s3_encryption_key_id": "arn:aws:kms:us-west-2:123:key/abc", + "aws_batch_role_arn": "arn:aws:iam::123:role/batch-role", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch-model" + ) + + assert credentials is not None + assert credentials["custom_llm_provider"] == "bedrock" + assert credentials["model"] == "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" + assert credentials["aws_region_name"] == "us-west-2" + assert credentials["s3_bucket_name"] == "my-batch-bucket" + assert credentials["s3_region_name"] == "us-east-1" + assert credentials["s3_encryption_key_id"] == "arn:aws:kms:us-west-2:123:key/abc" + assert credentials["aws_batch_role_arn"] == "arn:aws:iam::123:role/batch-role" + + def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict: return { "model_name": f"model_name_team-1_{model_id}", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0064b1a7d87..397d219b448 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26716,6 +26716,8 @@ export interface components { auto_router_max_input_chars?: number | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; + /** Aws Batch Role Arn */ + aws_batch_role_arn?: string | null; /** Aws Bedrock Project Id */ aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */ @@ -26949,6 +26951,8 @@ export interface components { s3_bucket_name?: string | null; /** S3 Encryption Key Id */ s3_encryption_key_id?: string | null; + /** S3 Region Name */ + s3_region_name?: string | null; /** Search Context Cost Per Query */ search_context_cost_per_query?: { [key: string]: unknown; @@ -35316,6 +35320,8 @@ export interface components { auto_router_max_input_chars?: number | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; + /** Aws Batch Role Arn */ + aws_batch_role_arn?: string | null; /** Aws Bedrock Project Id */ aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */ @@ -35549,6 +35555,8 @@ export interface components { s3_bucket_name?: string | null; /** S3 Encryption Key Id */ s3_encryption_key_id?: string | null; + /** S3 Region Name */ + s3_region_name?: string | null; /** Search Context Cost Per Query */ search_context_cost_per_query?: { [key: string]: unknown; From 86890654c5b96bdada40fc8e35b2812d8fd61284 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 5 Aug 2026 22:33:54 -0700 Subject: [PATCH 12/16] fix(proxy): include today's UTC bucket when a daily activity range ends at the caller's current day (#36051) * fix(proxy): include today's UTC bucket when a daily activity range ends at the caller's current day * fix(proxy): gate the current-UTC-day extension behind an opt-in param sent by the cost optimization dashboard * fix(ui): label cost optimization savings dates as UTC days --- .../common_daily_activity.py | 55 +++++++++++----- .../internal_user_endpoints.py | 8 +++ .../test_common_daily_activity.py | 62 +++++++++++++++++++ .../_components/UsageTab.test.tsx | 4 +- .../_components/UsageTab.tsx | 3 +- .../useDailyActivityRange.test.tsx | 4 +- .../_components/useDailyActivityRange.ts | 2 +- .../src/components/networking.tsx | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 9 files changed, 121 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 9af65b50c7f..7a30f6b799a 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -1,6 +1,6 @@ import asyncio from collections.abc import Awaitable, Callable, Mapping, Sequence -from datetime import datetime +from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import TYPE_CHECKING, Final, Protocol @@ -422,26 +422,46 @@ def _adjust_dates_for_timezone( start_date: str, end_date: str, timezone_offset_minutes: int | None, + include_current_utc_day: bool = False, + utc_now: datetime | None = None, ) -> tuple[str, str]: """ - Pass-through for the local date range; the timezone offset is intentionally ignored here. + Map a caller-local date range onto UTC bucket keys, extending only the live end. The aggregation table (e.g. LiteLLM_DailyUserSpend) stores spend in whole-UTC-day - buckets keyed on date as YYYY-MM-DD. Any conversion from a local date range to a - UTC date range using only date arithmetic must round to whole UTC days, allowing up - to 24h of slop at each boundary. The previous implementation expanded the SQL range - by an extra full UTC day on whichever side the offset pointed, which pulled in 24h - of unrelated bucket data per boundary and produced approximately 100% over-counting - on single-day queries (e.g. IST May 29 returning UTC May 28 + UTC May 29 in full). + buckets keyed on date as YYYY-MM-DD. Any conversion of an interior local-day + boundary using only date arithmetic must round to whole UTC days, allowing up to + 24h of slop at each boundary. A previous implementation expanded the SQL range by + an extra full UTC day on whichever side the offset pointed, which pulled in 24h of + unrelated bucket data per boundary and produced approximately 100% over-counting on + single-day queries (e.g. IST May 29 returning UTC May 28 + UTC May 29 in full). Sums of single-day queries then exceeded the equivalent multi-day aggregate, which - is mathematically impossible. + is mathematically impossible. Historical dates therefore stay a pass-through: the + local date is the UTC bucket key, trading boundary slop for monotonic, additive + results. Hour-level buckets or pro-rata weighting would fix that properly; both + require data the current schema does not store. - Treating the local date as the UTC date trades a small one-time boundary slop for - correct, monotonic, additive results across single-day and multi-day queries. A - later fix can introduce hour-level buckets or pro-rata weighting on adjacent UTC - days; both require data the current schema does not store. + The end boundary is different when the range reaches the caller's current day. A + caller west of UTC asking for a range ending "today" is asking for data up to now, + but once UTC has rolled past their local midnight, everything they sent since then + sits in the next UTC bucket, which the pass-through excludes: a PT dashboard goes + stale every evening from 5pm until local midnight, showing $0 for anything that + only started accruing that evening. Extending such a range to today's UTC bucket + cannot over-count, because the only part of that bucket outside the caller's range + is the future, and the future is empty. ``timezone_offset_minutes`` follows the + JS ``Date.getTimezoneOffset`` convention: UTC minus local, positive west of UTC. + + The extension is strictly opt-in via ``include_current_utc_day`` so a consumer + whose axis or reconciliation expects the range to stop at the requested end date + keeps today's byte-for-byte behaviour; the cost optimization dashboard opts in. """ - return start_date, end_date + if not include_current_utc_day or timezone_offset_minutes is None: + return start_date, end_date + now: Final = utc_now if utc_now is not None else datetime.now(timezone.utc) + caller_local_today: Final = (now - timedelta(minutes=timezone_offset_minutes)).date().isoformat() + if end_date < caller_local_today: + return start_date, end_date + return start_date, max(end_date, now.date().isoformat()) def _build_where_conditions( @@ -454,10 +474,13 @@ def _build_where_conditions( api_key: str | list[str] | None, exclude_entity_ids: list[str] | None = None, timezone_offset_minutes: int | None = None, + include_current_utc_day: bool = False, ) -> dict[str, "_WhereValue"]: """Build prisma where clause for daily activity queries.""" # Adjust dates for timezone if provided - adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes) + adjusted_start, adjusted_end = _adjust_dates_for_timezone( + start_date, end_date, timezone_offset_minutes, include_current_utc_day + ) where_conditions: Final[dict[str, _WhereValue]] = { "date": { @@ -903,6 +926,7 @@ async def get_daily_activity( exclude_entity_ids: list[str] | None = None, metadata_metrics_func: Callable[[Sequence[DailySpendRecord]], SpendMetrics] | None = None, timezone_offset_minutes: int | None = None, + include_current_utc_day: bool = False, resolve_entity_metadata: Callable[[Sequence[DailySpendRecord]], Awaitable[dict[str, dict[str, object]]]] | None = None, ) -> SpendAnalyticsPaginatedResponse: @@ -936,6 +960,7 @@ async def get_daily_activity( api_key=api_key, exclude_entity_ids=exclude_entity_ids, timezone_offset_minutes=timezone_offset_minutes, + include_current_utc_day=include_current_utc_day, ) # Get total count for pagination diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index dcec33f1cb2..640a735c916 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -2650,6 +2650,13 @@ async def get_user_daily_activity( description="Timezone offset in minutes from UTC (e.g., 480 for PST). " "Matches JavaScript's Date.getTimezoneOffset() convention.", ), + include_current_utc_day: bool = fastapi.Query( + default=False, + description="When the range ends on the caller's current local day, extend it to " + "today's UTC bucket so spend written after the caller's local midnight (in UTC " + "terms) is included. Requires the timezone parameter. Historical ranges are " + "never extended.", + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> SpendAnalyticsPaginatedResponse: """ @@ -2711,6 +2718,7 @@ async def get_user_daily_activity( page=page, page_size=page_size, timezone_offset_minutes=timezone, + include_current_utc_day=include_current_utc_day, resolve_entity_metadata=lambda records: _resolve_user_email_metadata(prisma_client, records), ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index f2749be5d6e..469e0d340f0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1,6 +1,8 @@ import os import sys +from datetime import datetime, timezone from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -870,6 +872,66 @@ class TestAdjustDatesForTimezone: assert per_day_ends == days +class TestAdjustDatesForTimezoneLiveEnd: + """ + Regression tests for the stale-evening bug: a caller west of UTC whose range + ends on their local "today" was capped at that local date's UTC bucket, so + once UTC rolled past their local midnight (5pm PT), everything sent that + evening sat in the next UTC bucket and the dashboard reported $0 for it + until local midnight. A range that reaches the caller's current day and + opts in via include_current_utc_day must extend to today's UTC bucket; the + only part of that bucket outside the range is the future, which is empty, + so the extension cannot over-count. Callers that do not opt in keep the + pass-through byte for byte. + """ + + PT_EVENING_UTC: Final = datetime(2026, 8, 6, 4, 30, tzinfo=timezone.utc) + + def test_pt_evening_range_ending_today_extends_to_utc_today(self): + start, end = _adjust_dates_for_timezone( + "2026-07-06", "2026-08-05", 420, include_current_utc_day=True, utc_now=self.PT_EVENING_UTC + ) + assert (start, end) == ("2026-07-06", "2026-08-06") + + def test_without_opt_in_live_range_keeps_pass_through(self): + start, end = _adjust_dates_for_timezone( + "2026-07-06", "2026-08-05", 420, utc_now=self.PT_EVENING_UTC + ) + assert (start, end) == ("2026-07-06", "2026-08-05") + + def test_pt_historical_range_is_untouched(self): + start, end = _adjust_dates_for_timezone( + "2026-07-01", "2026-08-04", 420, include_current_utc_day=True, utc_now=self.PT_EVENING_UTC + ) + assert (start, end) == ("2026-07-01", "2026-08-04") + + def test_east_of_utc_local_today_already_covers_utc_today(self): + ist_evening_utc: Final = datetime(2026, 8, 5, 17, 0, tzinfo=timezone.utc) + start, end = _adjust_dates_for_timezone( + "2026-07-07", "2026-08-06", -330, include_current_utc_day=True, utc_now=ist_evening_utc + ) + assert (start, end) == ("2026-07-07", "2026-08-06") + + def test_missing_offset_stays_pass_through_even_for_live_range(self): + start, end = _adjust_dates_for_timezone( + "2026-07-06", "2026-08-05", None, include_current_utc_day=True, utc_now=self.PT_EVENING_UTC + ) + assert (start, end) == ("2026-07-06", "2026-08-05") + + def test_utc_caller_range_ending_today_is_unchanged(self): + utc_noon: Final = datetime(2026, 8, 5, 12, 0, tzinfo=timezone.utc) + start, end = _adjust_dates_for_timezone( + "2026-07-06", "2026-08-05", 0, include_current_utc_day=True, utc_now=utc_noon + ) + assert (start, end) == ("2026-07-06", "2026-08-05") + + def test_future_end_date_extends_no_further_than_requested(self): + start, end = _adjust_dates_for_timezone( + "2026-07-06", "2026-08-09", 420, include_current_utc_day=True, utc_now=self.PT_EVENING_UTC + ) + assert (start, end) == ("2026-07-06", "2026-08-09") + + class TestBuildAggregatedSqlQuery: """ Asserts the SQL emitted by the aggregated query path stays anchored to the diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index fe3e792eeee..96d4644804b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -209,9 +209,9 @@ describe("UsageTab", () => { it("says what the line means and over what range", async () => { const { getByText, getByRole } = renderWith(twoDays()); - expect(getByText("Running total saved · Jul 1 – Jul 14")).toBeInTheDocument(); + expect(getByText("Running total saved · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); await userEvent.click(getByRole("tab", { name: "Per day" })); - expect(getByText("Saved per day · Jul 1 – Jul 14")).toBeInTheDocument(); + expect(getByText("Saved per day · Jul 1 – Jul 14 (UTC)")).toBeInTheDocument(); }); it("builds the per-driver donut from the range totals, not the running total", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx index b6287602210..bd9d4f3c873 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx @@ -141,7 +141,7 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { const rangeLabel = formatRangeLabel(startTime ?? undefined, endTime ?? undefined); const savingsSubtitle = [ accumulation === "cumulative" ? "Running total saved" : `Saved ${intervalLabel.toLowerCase()}`, - rangeLabel, + rangeLabel && `${rangeLabel} (UTC)`, ] .filter(Boolean) .join(" \u00b7 "); @@ -179,6 +179,7 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { return (
+ Spend is bucketed by UTC day
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx index 9fd27d80c37..e26a3629e8c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx @@ -22,13 +22,13 @@ describe("useDailyActivityRange", () => { it("queries every user's activity for an admin", () => { renderHook(() => useDailyActivityRange("test-token", "u1", "proxy_admin")); - expect(argsOfLastCall()).toEqual(["test-token", expect.any(Date), expect.any(Date), null]); + expect(argsOfLastCall()).toEqual(["test-token", expect.any(Date), expect.any(Date), null, true]); }); it("scopes the query to the caller for a non-admin", () => { renderHook(() => useDailyActivityRange("test-token", "u1", "internal_user")); - expect(argsOfLastCall()).toEqual(["test-token", expect.any(Date), expect.any(Date), "u1"]); + expect(argsOfLastCall()).toEqual(["test-token", expect.any(Date), expect.any(Date), "u1", true]); }); it("stays disabled until an access token is available", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts index 1c3f706726e..3a2a38c5955 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts @@ -35,7 +35,7 @@ export const useDailyActivityRange = ( const { data, loading, isFetchingMore } = usePaginatedDailyActivity({ fetchFn: userDailyActivityCall, - args: [accessToken, startTime, endTime, effectiveUserId], + args: [accessToken, startTime, endTime, effectiveUserId, true], enabled: !!accessToken && !!startTime && !!endTime, }); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 03106528a80..17a5ca37990 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1388,6 +1388,7 @@ export const userDailyActivityCall = async ( endTime: Date, page: number = 1, userId: string | null = null, + includeCurrentUtcDay: boolean = false, ) => { /** * Get daily user activity on proxy @@ -1400,6 +1401,7 @@ export const userDailyActivityCall = async ( page, extraQueryParams: { user_id: userId, + include_current_utc_day: includeCurrentUtcDay ? "true" : undefined, }, }); }; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0064b1a7d87..b67083ea90f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -53883,6 +53883,8 @@ export interface operations { page_size?: number; /** @description Timezone offset in minutes from UTC (e.g., 480 for PST). Matches JavaScript's Date.getTimezoneOffset() convention. */ timezone?: number | null; + /** @description When the range ends on the caller's current local day, extend it to today's UTC bucket so spend written after the caller's local midnight (in UTC terms) is included. Requires the timezone parameter. Historical ranges are never extended. */ + include_current_utc_day?: boolean; }; header?: never; path?: never; From 34fc8d2ee71f60b15cb097c90ca1649cc6377076 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 5 Aug 2026 22:34:55 -0700 Subject: [PATCH 13/16] fix: expired-miss share over all measured turns + cost-optimization tab labels (#36037) * fix(ui): make the expired-miss stat row a focusable tooltip trigger * fix: auto-router expired-miss percentage and cost-optimization tab labels - change expired-miss percentage denominator from return-to-tier misses to all measured turns (same_model + first_visit + return_to_tier). when auto-routers flip tiers rapidly within TTL, return-to-tier turns become hits and disappear from the miss count; the old metric reported only the rare failure population. the new metric contextualizes that population as a share of overall coverage - rename usage tab from 'Usage' to 'Overall' - rename auto-router-usage tab from 'Auto-Router Usage' to 'Auto-Router' - update component and unit tests to match new semantics --- .../AutoRouterBenchmarksTab.test.tsx | 30 +++++++++++++-- .../_components/AutoRouterBenchmarksTab.tsx | 38 ++++++++++--------- .../_components/CostOptimizationView.test.tsx | 10 ++--- .../_components/CostOptimizationView.tsx | 4 +- .../_components/autoRouterBenchmarks.test.ts | 21 ++++++++-- .../_components/autoRouterBenchmarks.ts | 6 +-- 6 files changed, 75 insertions(+), 34 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx index 51e9e125cb9..a5767383307 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx @@ -155,21 +155,45 @@ describe("AutoRouterBenchmarksTab", () => { expect(screen.getByText(/turns measured/)).toBeInTheDocument(); }); - it("recomputes the expired-miss share from the miss counts", () => { + it("computes the expired-miss share over every measured turn, not just return-to-tier misses", () => { mockHook({ data: response([group()]) }); renderTab(); expect(screen.getByText("Expired-miss")).toBeInTheDocument(); - expect(screen.getByText("27.1%")).toBeInTheDocument(); + expect(screen.getByText("2.3%")).toBeInTheDocument(); }); - it("hides the expired-miss row when every return turn hit", () => { + it("exposes the whole expired-miss row as a focusable tooltip trigger", () => { + mockHook({ data: response([group()]) }); + renderTab(); + + const trigger = screen.getByRole("button", { name: /Expired-miss/ }); + expect(trigger).toHaveTextContent("2.3%"); + }); + + it("shows a zero expired-miss share, rather than hiding the row, when every return turn hit", () => { const allHits = totals({ cache: cache({ return_to_tier: { turns: 381, hits: 381, hit_rate_pct: 100 }, return_misses_expired: 0 }), }); mockHook({ data: response([group(allHits)], allHits) }); renderTab(); + const trigger = screen.getByRole("button", { name: /Expired-miss/ }); + expect(trigger).toHaveTextContent("0.0%"); + }); + + it("hides the expired-miss row only when no turns were measured at all", () => { + const empty = { turns: 0, hits: 0, hit_rate_pct: 0 }; + const nothingMeasured = { + same_model: empty, + first_visit: empty, + return_to_tier: empty, + return_misses_expired: 0, + }; + const noTurns = totals({ cache: cache(nothingMeasured) }); + mockHook({ data: response([group(noTurns)], noTurns) }); + renderTab(); + expect(screen.queryByText("Expired-miss")).not.toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx index 2b71c36c597..ff0f52940b2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx @@ -183,23 +183,27 @@ const CachingCard: React.FC<{ cache: AutoRouterCacheStats }> = ({ cache }) => {

{pctLabel(cache.hit_rate_pct)}

{expiredMissPct === null ? null : ( -
- - - - Expired-miss -

- } - /> - - percentage of return-to-tier cache misses caused by cache expiring - - - -

{pctLabel(expiredMissPct)}

-
+ + + + } + > + + Expired-miss + + {pctLabel(expiredMissPct)} + + + share of all measured turns that missed cache because a return to an earlier tier came after its TTL + lapsed + + + )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx index 96ef75e8dd1..33c64ecf18a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx @@ -17,21 +17,21 @@ describe("CostOptimizationView", () => { it("renders the four cost-optimization tabs", () => { const { getByText } = renderView(); - expect(getByText("Usage")).toBeInTheDocument(); + expect(getByText("Overall")).toBeInTheDocument(); expect(getByText("Prompt Compression")).toBeInTheDocument(); expect(getByText("Prompt Caching")).toBeInTheDocument(); - expect(getByText("Auto-Router Usage")).toBeInTheDocument(); + expect(getByText("Auto-Router")).toBeInTheDocument(); }); - it("defaults to the Usage tab and switches the active tab on click", () => { + it("defaults to the Overall tab and switches the active tab on click", () => { const { getByRole } = renderView(); - expect(getByRole("tab", { name: "Usage" })).toHaveAttribute("aria-selected", "true"); + expect(getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "true"); expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "false"); fireEvent.click(getByRole("tab", { name: "Prompt Compression" })); - expect(getByRole("tab", { name: "Usage" })).toHaveAttribute("aria-selected", "false"); + expect(getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "false"); expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "true"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx index 0d986fdcaa8..6af1e8d0441 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx @@ -22,7 +22,7 @@ const CostOptimizationView: React.FC = ({ accessToken const items = [ { key: "usage", - label: "Usage", + label: "Overall", children: , }, { @@ -37,7 +37,7 @@ const CostOptimizationView: React.FC = ({ accessToken }, { key: "autorouter-usage", - label: "Auto-Router Usage", + label: "Auto-Router", children: , }, ]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts index 201a71ae4be..57c059b5524 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.test.ts @@ -131,12 +131,25 @@ describe("bucketRows", () => { }); describe("expiredMissShare", () => { - it("recomputes the expired share from the miss counts", () => { - expect(expiredMissShare(cache())).toBeCloseTo((100 * 19) / 70); + it("computes the expired share over every measured turn, not just return-to-tier misses", () => { + expect(expiredMissShare(cache())).toBeCloseTo((100 * 19) / 818); }); - it("is absent when every return turn hit", () => { - expect(expiredMissShare(cache({ return_to_tier: { turns: 10, hits: 10, hit_rate_pct: 100 } }))).toBeNull(); + it("is zero, not absent, when every return turn hit", () => { + expect( + expiredMissShare(cache({ return_to_tier: { turns: 10, hits: 10, hit_rate_pct: 100 }, return_misses_expired: 0 })), + ).toBe(0); + }); + + it("is absent only when no turns were measured at all", () => { + const empty = { turns: 0, hits: 0, hit_rate_pct: 0 }; + const nothingMeasured = { + same_model: empty, + first_visit: empty, + return_to_tier: empty, + return_misses_expired: 0, + }; + expect(expiredMissShare(cache(nothingMeasured))).toBeNull(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.ts index 8b6a4fa4105..00793548278 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks.ts @@ -91,9 +91,9 @@ export const bucketRows = (cache: AutoRouterCacheStats): BucketRow[] => { }; export const expiredMissShare = (cache: AutoRouterCacheStats): number | null => { - const misses = cache.return_to_tier.turns - cache.return_to_tier.hits; - if (misses <= 0) return null; - return (100 * cache.return_misses_expired) / misses; + const total = bucketTurnsTotal(cache); + if (total <= 0) return null; + return (100 * cache.return_misses_expired) / total; }; export const pctLabel = (value: number, digits: number = 1): string => `${value.toFixed(digits)}%`; From 7d745521bf94331cc5e0ec069e9c2f7d8d443303 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:46:40 -0700 Subject: [PATCH 14/16] fix(guardrails): merge synthesized tools under scan_only_tool_results and reject role-filtered no-op combos at init --- litellm/integrations/custom_guardrail.py | 10 +++ .../chat/guardrail_translation/handler.py | 14 +++- .../base_llm/guardrail_translation/utils.py | 43 +++++++++++- .../chat/guardrail_translation/handler.py | 14 +++- .../guardrail_hooks/bedrock_guardrails.py | 3 + .../panw_prisma_airs/panw_prisma_airs.py | 3 + .../proxy/guardrails/guardrail_registry.py | 12 ++++ .../test_anthropic_guardrail_handler.py | 8 ++- .../test_openai_guardrail_handler.py | 68 +++++++++++++++++-- .../guardrails/test_guardrail_registry.py | 57 ++++++++++++++++ 10 files changed, 216 insertions(+), 16 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 20f3aa430e9..7c8c9aeb248 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -714,6 +714,16 @@ class CustomGuardrail(CustomLogger): return result + def supports_scan_only_tool_results(self) -> bool: + """Whether this guardrail can scan tool-result content. + + Guardrails whose own role filtering only ever scans human-authored + messages override this to return False, so configuring them with + ``scan_only_tool_results`` is rejected at initialization instead of + silently scanning nothing on every request. + """ + return True + def should_run_guardrail( self, data, diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 184e0f6a343..aef01765e6e 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -26,10 +26,12 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im ) from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.llms.base_llm.guardrail_translation.utils import ( + anthropic_tool_name, effective_scan_only_tool_results_for_guardrail, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, merge_guardrailed_scoped_messages, + merge_returned_tools_into_request_tools, scoped_structured_message_indices, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( @@ -387,7 +389,7 @@ class AnthropicMessagesHandler(BaseTranslation): guardrailed_texts: Final = guardrailed_inputs.get("texts", []) guardrailed_tools: Final = guardrailed_inputs.get("tools") - if guardrailed_tools is not None and not scan_only_tool_results: + if guardrailed_tools is not None: # Convert tools back from OpenAI format to Anthropic format anthropic_config: Final = AnthropicConfig() anthropic_tools: Final[list[AllAnthropicToolsValues]] = [] @@ -396,7 +398,15 @@ class AnthropicMessagesHandler(BaseTranslation): if converted_tool is not None: anthropic_tools.append(converted_tool) # Note: MCP servers are handled separately in the main transformation - data["tools"] = anthropic_tools + data["tools"] = ( + merge_returned_tools_into_request_tools( + request_tools=data.get("tools"), + returned_tools=anthropic_tools, + tool_name=anthropic_tool_name, + ) + if scan_only_tool_results + else anthropic_tools + ) guardrailed_structured_messages: Final = guardrailed_inputs.get("structured_messages") if ( diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 432ac64b456..bdfe15ca9a6 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -1,8 +1,8 @@ from __future__ import annotations import json -from collections.abc import Iterator, Sequence -from typing import Any, Final +from collections.abc import Callable, Iterator, Sequence +from typing import Any, Final, TypeVar from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues @@ -167,6 +167,45 @@ def scoped_structured_message_indices( ) +ToolT = TypeVar("ToolT") + + +def openai_tool_name(tool: object) -> str | None: + if not isinstance(tool, dict): + return None + function: Final = tool.get("function") + if isinstance(function, dict): + function_name: Final = function.get("name") + return function_name if isinstance(function_name, str) else None + flat_name: Final = tool.get("name") + return flat_name if isinstance(flat_name, str) else None + + +def anthropic_tool_name(tool: object) -> str | None: + name: Final = tool.get("name") if isinstance(tool, dict) else None + return name if isinstance(name, str) else None + + +def merge_returned_tools_into_request_tools( + request_tools: Sequence[ToolT] | None, + returned_tools: Sequence[ToolT], + tool_name: Callable[[ToolT], str | None], +) -> list[ToolT]: + """Union of the request's tools and guardrail-returned tools, keyed by name. + + Under ``scan_only_tool_results`` the guardrail never saw the request's + tools, so a returned list can neither replace them (it would drop every + user-defined function) nor be discarded (it may carry a tool the guardrail + synthesized and told the model to call, like Compresr's retrieve tool). + Keep every request tool and append only returned tools whose names aren't + already taken. + """ + originals: Final = tuple(request_tools or ()) + taken_names: Final = frozenset(name for tool in originals if (name := tool_name(tool)) is not None) + additions: Final = tuple(tool for tool in returned_tools if tool_name(tool) not in taken_names) + return [*originals, *additions] + + def merge_guardrailed_scoped_messages( full_messages: Sequence[AllMessageValues], scoped_indices: Sequence[int], diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index dc2a06d67fc..4f0f69866ad 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -27,6 +27,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, merge_guardrailed_scoped_messages, + merge_returned_tools_into_request_tools, + openai_tool_name, role_out_of_guardrail_scope, scoped_structured_message_indices, ) @@ -143,8 +145,16 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrailed_texts: Final = guardrailed_inputs.get("texts", []) guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", []) guardrailed_tools: Final = guardrailed_inputs.get("tools") - if guardrailed_tools is not None and not scan_only_tool_results: - data["tools"] = guardrailed_tools + if guardrailed_tools is not None: + data["tools"] = ( + merge_returned_tools_into_request_tools( + request_tools=tools, + returned_tools=guardrailed_tools, + tool_name=openai_tool_name, + ) + if scan_only_tool_results + else guardrailed_tools + ) guardrailed_structured_messages: Final = guardrailed_inputs.get("structured_messages") if ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 8193069fd82..e9e729fb118 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -405,6 +405,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): grounding.append(block) return grounding + def supports_scan_only_tool_results(self) -> bool: + return self.experimental_use_latest_role_message_only is not True + def _prepare_guardrail_messages_for_role( self, messages: list[AllMessageValues] | None, 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 a96a0070eef..13ced0ac06c 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 @@ -1564,6 +1564,9 @@ class PanwPrismaAirsHandler(CustomGuardrail): return scannable + def supports_scan_only_tool_results(self) -> bool: + return False + @staticmethod def _get_scannable_text_indices( texts: list[str], diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index e9e61283c1a..15e884c939c 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -14,6 +14,9 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_scan_only_tool_results_for_guardrail, +) from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, ) @@ -493,6 +496,15 @@ class InMemoryGuardrailHandler: "scan_only_tool_results", ): setattr(custom_guardrail_callback, scoping_param, getattr(litellm_params, scoping_param, None)) + if ( + effective_scan_only_tool_results_for_guardrail(custom_guardrail_callback) + and not custom_guardrail_callback.supports_scan_only_tool_results() + ): + raise ValueError( + f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results is enabled, but this " + "guardrail's role filtering never scans tool results, so no request content would ever " + "be scanned. Remove scan_only_tool_results or the guardrail's role-filtering option." + ) configured_run_in_parallel: Final = getattr(litellm_params, "run_in_parallel", None) if configured_run_in_parallel is not None: custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel) 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 c7dedff0663..c7a30f7f954 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 @@ -883,7 +883,7 @@ class TestAnthropicMessagesScanOnlyToolResults: assert data["messages"][2]["content"][0]["text"] == "sibling POISON text" @pytest.mark.asyncio - async def test_guardrail_synthesized_tools_never_replace_scoped_out_request_tools(self): + async def test_guardrail_synthesized_tools_are_appended_without_replacing_request_tools(self): handler = AnthropicMessagesHandler() guardrail = ToolAppendingGuardrail(guardrail_name="tool-appending") guardrail.scan_only_tool_results = True @@ -912,9 +912,11 @@ class TestAnthropicMessagesScanOnlyToolResults: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - assert data["tools"] == original_tools, ( - "tools the guardrail synthesized without seeing the request's tools must not replace them" + assert [t["name"] for t in data["tools"]] == ["get_weather", "injected_tool"], ( + "a tool the guardrail synthesized must reach the model, converted to Anthropic format, " + "without the request's own tools being replaced or dropped" ) + assert data["tools"][0] == original_tools[0] @pytest.mark.asyncio async def test_guardrail_is_not_called_when_the_request_has_no_tool_results(self): 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 269afef69cd..deabee12497 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 @@ -1278,6 +1278,35 @@ class ToolSynthesizingGuardrail(CustomGuardrail): return inputs +class ToolNameCollidingGuardrail(CustomGuardrail): + """Returns a tool reusing a request tool's name plus a genuinely new tool.""" + + def __init__(self): + super().__init__(guardrail_name="tool-name-colliding") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + inputs["tools"] = [ + { + "type": "function", + "function": { + "name": "read_file", + "parameters": {"type": "object", "properties": {"hijacked": {"type": "string"}}}, + }, + }, + { + "type": "function", + "function": {"name": "injected_retrieve", "parameters": {"type": "object", "properties": {}}}, + }, + ] + return inputs + + class TestScanOnlyToolResults: def _bedrock_guardrail(self): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail @@ -1375,7 +1404,9 @@ class TestScanOnlyToolResults: @pytest.mark.parametrize("scan_only_tool_results", [True, False]) @pytest.mark.asyncio - async def test_guardrail_synthesized_tools_never_replace_scoped_out_request_tools(self, scan_only_tool_results): + async def test_guardrail_synthesized_tools_are_appended_without_replacing_request_tools( + self, scan_only_tool_results + ): handler = OpenAIChatCompletionsHandler() guardrail = ToolSynthesizingGuardrail() guardrail.scan_only_tool_results = scan_only_tool_results @@ -1395,12 +1426,35 @@ class TestScanOnlyToolResults: await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) - if scan_only_tool_results: - assert data["tools"] == original_tools, ( - "tools the guardrail synthesized without seeing the request's tools must not replace them" - ) - else: - assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"] + assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"], ( + "a tool the guardrail synthesized (like a recovery/retrieve tool) must reach the model " + "without the request's own tools being replaced or dropped" + ) + assert data["tools"][0] == original_tools[0] + + @pytest.mark.asyncio + async def test_returned_tool_name_collisions_keep_the_request_schema(self): + handler = OpenAIChatCompletionsHandler() + guardrail = ToolNameCollidingGuardrail() + guardrail.scan_only_tool_results = True + original_read_file = { + "type": "function", + "function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}}, + } + data = { + "messages": [ + {"role": "user", "content": "read the report"}, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}, + ], + "tools": [original_read_file], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"] + assert data["tools"][0] == original_read_file, ( + "a returned tool reusing a request tool's name must not replace the request's schema" + ) @pytest.mark.asyncio async def test_structured_write_back_keeps_out_of_scope_messages(self): diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 6bd109f0f95..3053664ef27 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -558,3 +558,60 @@ def test_reinitialized_judge_guardrail_uses_lazy_router_provider(): finally: for cb_list, snapshot in zip(lists, snapshots): cb_list[:] = snapshot + + +class TestScanOnlyToolResultsInitRefusal: + """A guardrail whose role filtering never scans tool results must be rejected at + initialization when configured with scan_only_tool_results, instead of booting a + proxy that silently scans nothing on every request.""" + + def _initialize(self, name: str, params: dict): + lists = _all_callback_lists() + snapshots = [list(cb_list) for cb_list in lists] + try: + return InMemoryGuardrailHandler().initialize_guardrail( + guardrail={"guardrail_name": name, "litellm_params": params}, + ) + finally: + for cb_list, snapshot in zip(lists, snapshots): + cb_list[:] = snapshot + + def test_panw_prisma_airs_with_scan_only_tool_results_is_rejected(self): + with pytest.raises(ValueError, match="never scans tool results"): + self._initialize( + "panw-scan-only-combo", + { + "guardrail": "panw_prisma_airs", + "mode": "pre_call", + "api_key": "test-key", + "profile_name": "test-profile", + "scan_only_tool_results": True, + }, + ) + + def test_bedrock_latest_role_with_scan_only_tool_results_is_rejected(self): + with pytest.raises(ValueError, match="never scans tool results"): + self._initialize( + "bedrock-latest-role-scan-only-combo", + { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "gr-1", + "guardrailVersion": "1", + "experimental_use_latest_role_message_only": True, + "scan_only_tool_results": True, + }, + ) + + def test_bedrock_without_latest_role_accepts_scan_only_tool_results(self): + result = self._initialize( + "bedrock-scan-only-ok", + { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "gr-1", + "guardrailVersion": "1", + "scan_only_tool_results": True, + }, + ) + assert result is not None From 28ff7f3f0b68d77428f175e54099b9b590226081 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:52:16 -0700 Subject: [PATCH 15/16] fix(guardrails): scan function-role results and dedupe returned tools Under scan_only_tool_results, legacy OpenAI function-role messages now count as tool results, and duplicate names among guardrail-returned tools keep only the first occurrence. CustomGuardrail.structured_messages_cover_full_request lets CrowdStrike AIDR declare that its writeback already rebuilds the whole conversation, so handlers install it as-is instead of merging it into the full message list a second time and duplicating out-of-scope rows. Lint budget ceilings ratchet down to match the tree --- basedpyright-code-budget.json | 18 ++--- litellm/integrations/custom_guardrail.py | 13 ++++ .../chat/guardrail_translation/handler.py | 4 +- .../base_llm/guardrail_translation/utils.py | 11 ++- .../chat/guardrail_translation/handler.py | 12 ++- .../crowdstrike_aidr/crowdstrike_aidr.py | 4 + ruff-strict-budget.json | 2 +- .../test_openai_guardrail_handler.py | 78 +++++++++++++++++++ type-discipline-budget.json | 6 +- 9 files changed, 127 insertions(+), 21 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 8a5c78c1f6c..d98e6c6c911 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -9,10 +9,10 @@ "limit": 329 }, "reportAttributeAccessIssue": { - "limit": 516 + "limit": 514 }, "reportCallIssue": { - "limit": 123 + "limit": 117 }, "reportConstantRedefinition": { "limit": 40 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 9226 + "limit": 9225 }, "reportFunctionMemberAccess": { "limit": 7 @@ -99,25 +99,25 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45242 + "limit": 45145 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40339 + "limit": 39881 }, "reportUnknownParameterType": { - "limit": 20293 + "limit": 20258 }, "reportUnknownVariableType": { - "limit": 31796 + "limit": 31429 }, "reportUnnecessaryCast": { "limit": 122 }, "reportUnnecessaryComparison": { - "limit": 702 + "limit": 701 }, "reportUnnecessaryContains": { "limit": 5 @@ -126,7 +126,7 @@ "limit": 864 }, "reportUntypedBaseClass": { - "limit": 72 + "limit": 0 }, "reportUntypedFunctionDecorator": { "limit": 33 diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 7c8c9aeb248..2e91e082bd4 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -724,6 +724,19 @@ class CustomGuardrail(CustomLogger): """ return True + def structured_messages_cover_full_request(self) -> bool: + """Whether returned ``structured_messages`` span the whole request. + + Translation handlers hand guardrails only the in-scope subset of the + conversation and merge a returned ``structured_messages`` list back + into the full request. A guardrail that already rebuilds the complete + conversation itself (like CrowdStrike AIDR with its skip filters + active) overrides this to return True so the handler installs the + returned list as-is instead of merging it a second time, which would + duplicate the out-of-scope messages. + """ + return False + def should_run_guardrail( self, data, diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index aef01765e6e..88db9fae912 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -415,7 +415,9 @@ class AnthropicMessagesHandler(BaseTranslation): ): self._write_back_structured_messages( data, - merge_guardrailed_scoped_messages( + guardrailed_structured_messages + if guardrail_to_apply.structured_messages_cover_full_request() + else merge_guardrailed_scoped_messages( full_messages=full_structured_messages, scoped_indices=scoped_message_indices, guardrailed_scoped=guardrailed_structured_messages, diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index bdfe15ca9a6..f1ddf21cd3c 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -145,7 +145,7 @@ def role_out_of_guardrail_scope( return True if skip_tool_message and role == "tool": return True - return scan_only_tool_results and role != "tool" + return scan_only_tool_results and role not in ("tool", "function") def scoped_structured_message_indices( @@ -198,11 +198,16 @@ def merge_returned_tools_into_request_tools( user-defined function) nor be discarded (it may carry a tool the guardrail synthesized and told the model to call, like Compresr's retrieve tool). Keep every request tool and append only returned tools whose names aren't - already taken. + already taken by a request tool or an earlier returned tool. """ originals: Final = tuple(request_tools or ()) taken_names: Final = frozenset(name for tool in originals if (name := tool_name(tool)) is not None) - additions: Final = tuple(tool for tool in returned_tools if tool_name(tool) not in taken_names) + additions: Final = tuple( + tool + for index, tool in enumerate(returned_tools) + if (name := tool_name(tool)) not in taken_names + and (name is None or all(tool_name(earlier) != name for earlier in returned_tools[:index])) + ) return [*originals, *additions] diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 4f0f69866ad..e411dc497fc 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -161,10 +161,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrailed_structured_messages is not None and guardrailed_structured_messages is not original_structured_messages ): - data["messages"] = merge_guardrailed_scoped_messages( - full_messages=structured_messages or [], - scoped_indices=scoped_message_indices, - guardrailed_scoped=guardrailed_structured_messages, + data["messages"] = ( + guardrailed_structured_messages + if guardrail_to_apply.structured_messages_cover_full_request() + else merge_guardrailed_scoped_messages( + full_messages=structured_messages or [], + scoped_indices=scoped_message_indices, + guardrailed_scoped=guardrailed_structured_messages, + ) ) else: # Step 3: Map guardrail responses back to original message structure diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 15ddd5e3458..b1bf9159607 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -362,6 +362,10 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): tail: Final = guard_output.messages[-num_assistant_messages:] if num_assistant_messages > 0 else [] return [_extract_text_from_message(msg) for msg in tail] + @override + def structured_messages_cover_full_request(self) -> bool: + return effective_skip_system_message_for_guardrail(self) or effective_skip_tool_message_for_guardrail(self) + def _writeback_messages( self, structured_messages: list[AllMessageValues], diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index ea20ac97e07..65c98f6aab3 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -42,7 +42,7 @@ "limit": 81 }, "B010": { - "limit": 192 + "limit": 190 }, "B018": { "limit": 2 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 deabee12497..2e75f29b1c5 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 @@ -1307,6 +1307,38 @@ class ToolNameCollidingGuardrail(CustomGuardrail): return inputs +class DuplicateToolReturningGuardrail(CustomGuardrail): + """Returns the same synthesized tool name twice, second copy with a different schema.""" + + def __init__(self): + super().__init__(guardrail_name="duplicate-tool-returning") + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + inputs["tools"] = [ + { + "type": "function", + "function": { + "name": "injected_retrieve", + "parameters": {"type": "object", "properties": {"first": {"type": "string"}}}, + }, + }, + { + "type": "function", + "function": { + "name": "injected_retrieve", + "parameters": {"type": "object", "properties": {"second": {"type": "string"}}}, + }, + }, + ] + return inputs + + class TestScanOnlyToolResults: def _bedrock_guardrail(self): from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail @@ -1351,6 +1383,28 @@ class TestScanOnlyToolResults: scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] assert scanned == ["TOOL-RESULT-scanned"] + @pytest.mark.asyncio + async def test_legacy_function_role_results_are_scanned(self): + from unittest.mock import AsyncMock, patch + + handler = OpenAIChatCompletionsHandler() + guardrail = self._bedrock_guardrail() + data = { + "messages": [ + {"role": "user", "content": "USER-PROMPT-not-scanned"}, + {"role": "function", "name": "read_file", "content": "FUNCTION-RESULT-scanned"}, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT-scanned"}, + ] + } + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + assert mock_api.call_count == 1 + scanned = [m["content"] for m in mock_api.call_args.kwargs["messages"]] + assert scanned == ["FUNCTION-RESULT-scanned", "TOOL-RESULT-scanned"], ( + "a tool result sent with the legacy function role must not bypass the scoped scan" + ) + @pytest.mark.parametrize("flag_value", [None, "false", 0, object()]) @pytest.mark.asyncio async def test_scope_narrows_only_when_the_flag_is_actually_true(self, flag_value): @@ -1456,6 +1510,30 @@ class TestScanOnlyToolResults: "a returned tool reusing a request tool's name must not replace the request's schema" ) + @pytest.mark.asyncio + async def test_duplicate_returned_tool_names_keep_only_the_first(self): + handler = OpenAIChatCompletionsHandler() + guardrail = DuplicateToolReturningGuardrail() + guardrail.scan_only_tool_results = True + original_read_file = { + "type": "function", + "function": {"name": "read_file", "parameters": {"type": "object", "properties": {}}}, + } + data = { + "messages": [ + {"role": "user", "content": "read the report"}, + {"role": "tool", "tool_call_id": "call_1", "content": "TOOL-RESULT"}, + ], + "tools": [original_read_file], + } + + await handler.process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert [t["function"]["name"] for t in data["tools"]] == ["read_file", "injected_retrieve"], ( + "two returned tools sharing a name must not both be forwarded to the provider" + ) + assert data["tools"][1]["function"]["parameters"]["properties"] == {"first": {"type": "string"}} + @pytest.mark.asyncio async def test_structured_write_back_keeps_out_of_scope_messages(self): handler = OpenAIChatCompletionsHandler() diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 37964c27657..8064e63f1aa 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23337 + "limit": 23332 }, "LIT002": { "limit": 27213 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1092 + "limit": 1091 }, "LIT007": { "limit": 0 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16796 + "limit": 16792 }, "LIT011": { "limit": 5602 From 14d4897e55a8223ee3b2815b0b4e038caa3c0f61 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 01:58:06 -0700 Subject: [PATCH 16/16] fix(guardrails): refuse scan_only_tool_results combos that scan nothing Prompt Security drops tool and function rows unless check_tool_results is on, so it now reports scan-only support from that setting and the registry refuses the pairing at boot. Pairing scan_only_tool_results with skip_tool_message_in_guardrail excludes every message, so guardrail initialization now rejects that combination too. --- .../prompt_security/prompt_security.py | 3 ++ .../proxy/guardrails/guardrail_registry.py | 15 +++++-- .../guardrails/test_guardrail_registry.py | 42 +++++++++++++++++++ 3 files changed, 56 insertions(+), 4 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 743ad888949..1a2c46f306c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -74,6 +74,9 @@ class PromptSecurityGuardrail(CustomGuardrail): super().__init__(**kwargs) + def supports_scan_only_tool_results(self) -> bool: + return self.check_tool_results + @log_guardrail_information async def apply_guardrail( self, diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 15e884c939c..9f70ed63dcb 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -16,6 +16,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.guardrail_translation.utils import ( effective_scan_only_tool_results_for_guardrail, + effective_skip_tool_message_for_guardrail, ) from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrail, @@ -496,15 +497,21 @@ class InMemoryGuardrailHandler: "scan_only_tool_results", ): setattr(custom_guardrail_callback, scoping_param, getattr(litellm_params, scoping_param, None)) - if ( - effective_scan_only_tool_results_for_guardrail(custom_guardrail_callback) - and not custom_guardrail_callback.supports_scan_only_tool_results() - ): + scan_only_tool_results_enabled: Final = effective_scan_only_tool_results_for_guardrail( + custom_guardrail_callback + ) + if scan_only_tool_results_enabled and not custom_guardrail_callback.supports_scan_only_tool_results(): raise ValueError( f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results is enabled, but this " "guardrail's role filtering never scans tool results, so no request content would ever " "be scanned. Remove scan_only_tool_results or the guardrail's role-filtering option." ) + if scan_only_tool_results_enabled and effective_skip_tool_message_for_guardrail(custom_guardrail_callback): + raise ValueError( + f"Guardrail {guardrail['guardrail_name']}: scan_only_tool_results and " + "skip_tool_message_in_guardrail are enabled together, which excludes every message from " + "scanning, so no request content would ever be scanned. Remove one of the two." + ) configured_run_in_parallel: Final = getattr(litellm_params, "run_in_parallel", None) if configured_run_in_parallel is not None: custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 3053664ef27..729dbce6b9a 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -615,3 +615,45 @@ class TestScanOnlyToolResultsInitRefusal: }, ) assert result is not None + + def test_prompt_security_default_tool_filtering_rejects_scan_only_tool_results(self, monkeypatch): + monkeypatch.delenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", raising=False) + with pytest.raises(ValueError, match="never scans tool results"): + self._initialize( + "prompt-security-scan-only-combo", + { + "guardrail": "prompt_security", + "mode": "pre_call", + "api_key": "test-key", + "api_base": "https://ps.example.com", + "scan_only_tool_results": True, + }, + ) + + def test_prompt_security_check_tool_results_accepts_scan_only_tool_results(self, monkeypatch): + monkeypatch.setenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", "true") + result = self._initialize( + "prompt-security-scan-only-ok", + { + "guardrail": "prompt_security", + "mode": "pre_call", + "api_key": "test-key", + "api_base": "https://ps.example.com", + "scan_only_tool_results": True, + }, + ) + assert result is not None + + def test_skip_tool_message_with_scan_only_tool_results_is_rejected(self): + with pytest.raises(ValueError, match="skip_tool_message_in_guardrail are enabled together"): + self._initialize( + "bedrock-skip-tool-scan-only-combo", + { + "guardrail": "bedrock", + "mode": "pre_call", + "guardrailIdentifier": "gr-1", + "guardrailVersion": "1", + "skip_tool_message_in_guardrail": True, + "scan_only_tool_results": True, + }, + )