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

@@ -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,

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

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

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

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"}) == {

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()

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"}',
+    }

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="{}")