From 8921bd1aa81f073a547e4384060b4891da3dfddf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 02:50:22 -0800 Subject: [PATCH] fix(guardrails): keep hosted web call copies in sync on request masks and denied server calls, and accept null tool_calls in Lakera --- litellm/llms/anthropic/common_utils.py | 53 ++++++++++++++++++- litellm/proxy/guardrails/_content_utils.py | 14 ++++- .../guardrails/guardrail_hooks/lakera_ai.py | 2 +- .../guardrail_hooks/tool_permission.py | 20 ++++--- .../test_secret_detection.py | 28 ++++++++++ .../test_openai_guardrail_handler.py | 33 ++++++++++++ .../guardrail_hooks/test_lakera_ai.py | 41 ++++++++++++++ .../guardrail_hooks/test_tool_permission.py | 24 +++++++++ 8 files changed, 203 insertions(+), 12 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 709164c6277..6aaeb21bafb 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -16,6 +16,7 @@ from urllib.parse import quote, urlparse import httpx from pydantic import BeforeValidator, ConfigDict, Field, JsonValue, StrictBool, TypeAdapter, ValidationError +from pydantic_core import to_jsonable_python from typing_extensions import ReadOnly, TypedDict import litellm @@ -2405,8 +2406,12 @@ def write_anthropic_server_tool_calls( def write_request_server_tool_calls(message: object, rewritten: tuple[ChatCompletionAssistantToolCall, ...]) -> None: + original: Final = anthropic_server_tool_call_params(message) + if rewritten == original: + return + write_request_hosted_web_call_arguments(message, changed_server_tool_call_arguments(original, rewritten)) fields: Final = validated_object_mapping(message_field(message, "provider_specific_fields")) or {} - if rewritten == anthropic_server_tool_call_params(message) or not isinstance(message, dict): + if not isinstance(message, dict): return message["provider_specific_fields"] = { # rebind-ok: a guardrail rewrite must land on the request it scanned **fields, @@ -2414,6 +2419,20 @@ def write_request_server_tool_calls(message: object, rewritten: tuple[ChatComple } +def write_request_hosted_web_call_arguments(message: object, arguments_by_tool_id: Mapping[str, str]) -> None: + fields: Final = validated_object_mapping(message_field(message, "provider_specific_fields")) + if not arguments_by_tool_id or not isinstance(message, dict) or fields is None: + return + if fields.get("web_search_calls") is None: + return + message["provider_specific_fields"] = { # rebind-ok: a rewritten call's hosted copy must not keep its old input + **fields, + "web_search_calls": to_jsonable_python( + _rewritten_web_search_calls(fields["web_search_calls"], arguments_by_tool_id), exclude_none=True + ), + } + + def _server_tool_call_id(call: ChatCompletionAssistantToolCall) -> str: return call["id"] or "" @@ -2467,6 +2486,38 @@ def _with_hosted_web_call_arguments(item: object, arguments_by_tool_id: Mapping[ return call.model_copy(update={"action": with_hosted_web_tool_arguments(call.action, arguments)}) +def without_anthropic_server_tool_calls(holder: Message, tool_ids: frozenset[str]) -> None: + fields: Final = holder.provider_specific_fields or {} + holder.provider_specific_fields = { # rebind-ok: a denied call must leave every copy of the message it was on + **fields, + "server_tool_calls": [ + call for call in anthropic_server_tool_call_params(holder) if _server_tool_call_id(call) not in tool_ids + ], + **( + {"web_search_calls": _web_search_calls_without(fields["web_search_calls"], tool_ids)} + if fields.get("web_search_calls") is not None + else {} + ), + } + + +def _web_search_calls_without(calls: object, tool_ids: frozenset[str]) -> object: + mapping: Final = validated_object_mapping(calls) + if mapping is not None: + return {key: call for key, call in mapping.items() if _hosted_web_call_tool_id(call) not in tool_ids} + items: Final = validated_object_list(calls) + if items is None: + return calls + return [call for call in items if _hosted_web_call_tool_id(call) not in tool_ids] + + +def _hosted_web_call_tool_id(item: object) -> str | None: + try: + return web_search_call_tool_id(_WEB_SEARCH_CALL_ADAPTER.validate_python(item).id) + except ValidationError: + return None + + def write_anthropic_server_tool_call_arguments(holder: Message | Delta, index: int, arguments: str) -> None: write_anthropic_server_tool_calls( holder, diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 2a1b70d6439..1a1670715dc 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -18,6 +18,11 @@ from typing import Final from pydantic import JsonValue, TypeAdapter from litellm.litellm_core_utils.core_helpers import iter_items +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + changed_server_tool_call_arguments, + write_request_hosted_web_call_arguments, +) from litellm.responses.litellm_completion_transformation.compaction_items import ( compaction_summary, input_items, @@ -413,8 +418,9 @@ def walk_user_text( if isinstance(function, dict) and "arguments" in function: function["arguments"] = _rewrite_arguments(function["arguments"]) - messages: Final = data.get("messages") - for message in _object_list(messages) or (): + messages: Final = _object_list(data.get("messages")) or () + server_tool_calls_before: Final = tuple(anthropic_server_tool_call_params(message) for message in messages) + for message in messages: if (skip_tool_message and _is_tool_message(message)) or not isinstance(message, dict): continue if "content" in message: @@ -423,6 +429,10 @@ def walk_user_text( provider_fields = message.get("provider_specific_fields") if isinstance(provider_fields, dict): _rewrite_call_arguments(provider_fields.get("server_tool_calls")) + for message, original in zip(messages, server_tool_calls_before, strict=True): + write_request_hosted_web_call_arguments( + message, changed_server_tool_call_arguments(original, anthropic_server_tool_call_params(message)) + ) if "system" in data: data["system"] = _rewrite_content(data["system"]) # rebind-ok: preserve the native hook's in-place contract diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index e94f7329885..9c8c9cbe96d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -167,7 +167,7 @@ class lakeraAI_Moderation(CustomGuardrail): tool_call_messages = [ *tool_call_messages, *anthropic_server_tool_call_params(message), - *message.get("tool_calls", ()), + *(message.get("tool_calls") or ()), ] if role == Role.SYSTEM.value: # we need this for later system_message = message diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 3577889ae0f..f2697c0aa11 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -24,7 +24,11 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.tool_fields import declared_function_name, qualified_member_name -from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params, anthropic_server_tool_calls +from litellm.llms.anthropic.common_utils import ( + anthropic_server_tool_call_params, + anthropic_server_tool_calls, + without_anthropic_server_tool_calls, +) from litellm.llms.base_llm.guardrail_translation.utils import anthropic_tool_names from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import ( @@ -1500,16 +1504,16 @@ class ToolPermissionGuardrail(CustomGuardrail): ) if any(server_error_results): error_messages.extend(result.content for result in server_error_results if result is not None) - choice.message.provider_specific_fields = { - **(choice.message.provider_specific_fields or {}), - "server_tool_calls": [ - server_tool_call + without_anthropic_server_tool_calls( + choice.message, + frozenset( + server_tool_call["id"] or "" for server_tool_call, result in zip( anthropic_server_tool_call_params(choice.message), server_error_results ) - if result is None - ], - } + if result is not None + ), + ) legacy_tool_call = self._legacy_function_call_to_tool_call( getattr(choice.message, "function_call", None), choice_index diff --git a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py index a9180f8dc63..3599c938c28 100644 --- a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py @@ -29,6 +29,7 @@ from litellm_enterprise.enterprise_callbacks.secret_detection import ( from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import hash_token +from litellm.types.responses.main import build_web_search_call AWS_KEY = "AKIAIOSFODNN7EXAMPLE" OPENAI_KEY = "sk-test-abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH" @@ -132,6 +133,33 @@ async def test_a_replayed_server_tool_call_cannot_forward_unredacted_credentials assert json.loads(rewritten) == {"query": "AWS_ACCESS_KEY_ID=[REDACTED]"} +@pytest.mark.asyncio +async def test_a_replayed_server_web_search_keeps_no_credentials_in_its_hosted_copy(): + query = f"AWS_ACCESS_KEY_ID={AWS_KEY}" + server_call = { + "id": "srvtoolu_1", + "type": "function", + "function": {"name": "web_search", "arguments": json.dumps({"query": query})}, + } + data = { + "messages": [ + { + "role": "assistant", + "content": "Searching.", + "provider_specific_fields": { + "server_tool_calls": [server_call], + "web_search_calls": [ + build_web_search_call("srvtoolu_1", {"query": query}, {"content": []}).model_dump() + ], + }, + } + ] + } + await _guardrail().async_pre_call_hook(UserAPIKeyAuth(), DualCache(), data, "acompletion") + hosted = json.loads(json.dumps(data["messages"][0]["provider_specific_fields"]["web_search_calls"])) + assert [call["action"]["query"] for call in hosted] == ["AWS_ACCESS_KEY_ID=[REDACTED]"] + + @pytest.mark.parametrize("password", ("hunter2\\brahms9x", "hunter2\tbrahms9x")) def test_tool_payload_redacts_keyword_values_that_need_json_escaping(password): assert _guardrail().redact_tool_payload({"password": password, "path": ".env"}) == { diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 697f20a6fa3..54832bb35cc 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -986,6 +986,39 @@ class TestOpenAIChatCompletionsHandlerToolCallsOutput: ("toolu_01Weather", "get_weather", '{"city": "Paris"}') ] + @pytest.mark.asyncio + async def test_input_guardrail_rewrite_of_a_replayed_server_web_call_reaches_its_hosted_copy(self): + handler: Final = OpenAIChatCompletionsHandler() + data: Final = { + "messages": [ + {"role": "user", "content": "look up jane"}, + { + "role": "assistant", + "content": "Searching.", + "provider_specific_fields": { + "server_tool_calls": [ + _server_tool_call("srvtoolu_01Search", "web_search", _SEARCH_INPUT), + _server_tool_call("srvtoolu_01Fetch", "web_fetch", _FETCH_INPUT), + ], + "web_search_calls": [ + build_web_search_call("srvtoolu_01Search", _SEARCH_INPUT, {"content": []}).model_dump(), + build_web_search_call( + "srvtoolu_01Fetch", _FETCH_INPUT, {"content": []}, tool_name="web_fetch" + ).model_dump(), + ], + }, + }, + ] + } + + await handler.process_input_messages(data, QueryRedactingGuardrail(guardrail_name="redact")) + + hosted: Final = json.loads(json.dumps(data["messages"][1]["provider_specific_fields"]["web_search_calls"])) + assert [(call["id"], call["action"].get("query") or call["action"].get("url")) for call in hosted] == [ + ("ws_srvtoolu_01Search", "[REDACTED]"), + ("ws_srvtoolu_01Fetch", "[REDACTED]"), + ] + @pytest.mark.asyncio async def test_output_guardrail_sees_a_response_that_only_carries_server_tool_calls(self): handler: Final = OpenAIChatCompletionsHandler() diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py index 6880d0ac491..cb7fb7898ab 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -55,3 +55,44 @@ async def test_prompt_check_appends_replayed_server_tool_call_arguments_to_syste "role": "system", "content": 'You are a shopping assistant. Function Input: {"url": "https://example.com"} {"order_id": "A-17"}', } + + +@pytest.mark.asyncio +async def test_prompt_check_scans_a_replayed_turn_that_only_has_server_tool_calls( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "add_function_to_prompt", False) + route: Final = respx_mock.post("https://api.lakera.ai/v1/prompt_injection").mock( + return_value=httpx.Response(200, json={"results": [{"flagged": False}]}) + ) + moderation: Final = lakeraAI_Moderation(api_key="test-key", moderation_check="pre_call") + + await moderation._check( + data={ + "messages": [ + {"role": "system", "content": "You are a shopping assistant."}, + { + "role": "assistant", + "content": "Searching.", + "tool_calls": None, + "provider_specific_fields": { + "server_tool_calls": [ + { + "id": "srvtoolu_01", + "type": "function", + "function": {"name": "web_search", "arguments": '{"query": "order A-17"}'}, + } + ] + }, + }, + ] + }, + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + + assert json.loads(route.calls.last.request.content)["input"][0] == { + "role": "system", + "content": 'You are a shopping assistant. Function Input: {"query": "order A-17"}', + } diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 5c481e405db..a46f1ddeed9 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -23,6 +23,7 @@ from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ( ToolPermissionGuardrail, ) from litellm.types.guardrails import GuardrailEventHooks, LitellmParams +from litellm.types.responses.main import build_web_search_call from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ) @@ -1448,6 +1449,29 @@ class TestToolPermissionGuardrail: assert [tool_call.id for tool_call in message.tool_calls] == [_ALLOWED_BASH_CALL["id"]] assert "Permission denied" in message.content + @pytest.mark.asyncio + async def test_rewrite_mode_removes_the_denied_server_tool_calls_hosted_copy(self): + rewriting: Final = ToolPermissionGuardrail( + guardrail_name="test-tool-permission", + rules=self.test_rules, + default_action="deny", + on_disallowed_action="rewrite", + ) + response: Final = self._server_search_response() + response.choices[0].message.provider_specific_fields = { + **response.choices[0].message.provider_specific_fields, + "web_search_calls": [ + build_web_search_call("srvtoolu_01Search", {"query": "weather in Paris"}, {"content": []}) + ], + } + + with patch.object(rewriting, "should_run_guardrail", return_value=True): + result: Final = await rewriting.async_post_call_success_hook( + data={}, user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert result.choices[0].message.provider_specific_fields["web_search_calls"] == [] + def test_modify_response_with_permission_errors_keeps_custom_tool_call_in_another_choice(self): denied_call = ChatCompletionMessageToolCall( id="call_bash", type="function", function=Function(name="Bash", arguments="{}")