From 8ea78482cc2c1043ce819ce174076f6cbfd40a94 Mon Sep 17 00:00:00 2001 From: Jesus Velarde Date: Fri, 6 Nov 2026 08:45:05 -0800 Subject: [PATCH] fix(guardrails): write rewritten tool_calls from the Generic Guardrail API back to the client The guardrail's response tool_calls were dropped at parse time and the /v1/messages output handler never applied them, so a post_call guardrail that masks tool-call arguments had no effect on any endpoint. The hook now parses tool_calls, matches them one to one with the calls it sent, keeps each call's own id and type, and forwards them. An empty array is no rewrite. A list of another length or a malformed entry is a guardrail error governed by fail_on_error. The Anthropic output handler writes the returned calls onto the tool_use blocks by position. --- .../chat/guardrail_translation/handler.py | 73 ++++- .../generic_guardrail_api.py | 16 +- .../guardrail_hooks/generic_guardrail_api.py | 61 +++- .../test_anthropic_guardrail_handler.py | 287 ++++++++++++++++-- .../test_generic_guardrail_api.py | 108 +++++++ 5 files changed, 512 insertions(+), 33 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 8f1a08d2964..6b84561af34 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -156,6 +156,12 @@ class _TextAttrBlock(Protocol): text: str +@runtime_checkable +class _ToolUseAttrBlock(Protocol): + name: str + input: Mapping[str, object] + + @runtime_checkable class _WritableMessage(Protocol): @overload @@ -443,13 +449,23 @@ def _rewritten_tool_use_input(arguments: str) -> Mapping[str, object] | None: def _write_back_tool_use( message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape, rewritten_input: Mapping[str, object] ) -> None: - content: Final = message.get("content", None) - block: Final = content[target.content_idx] if isinstance(content, list) else None - if not isinstance(block, dict): + blocks: Final = list_field(message, "content") + _write_back_tool_use_block( + blocks[target.content_idx] if target.content_idx < len(blocks) else None, shape, rewritten_input + ) + + +def _write_back_tool_use_block(block: object, shape: _ToolCallShape, rewritten_input: Mapping[str, object]) -> None: + if isinstance(block, dict): + block["input"] = rewritten_input # rebind-ok: the tool_use block is the client-visible object being rewritten + if shape.name is not None and shape.name != block.get("name"): + block["name"] = shape.name # rebind-ok: same in-place rewrite as the input above return - block["input"] = rewritten_input - if shape.name is not None and shape.name != block.get("name"): - block["name"] = shape.name + if not isinstance(block, _ToolUseAttrBlock): + return + block.input = rewritten_input # rebind-ok: same in-place rewrite for an object-shaped tool_use block + if shape.name is not None and shape.name != block.name: + block.name = shape.name # rebind-ok: same in-place rewrite as the input above def _streamed_query_strings( @@ -1520,6 +1536,7 @@ class AnthropicMessagesHandler(BaseTranslation): tool_calls_to_check, response, ) + pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check) guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -1541,6 +1558,13 @@ class AnthropicMessagesHandler(BaseTranslation): self._write_back_web_search_strings( response_content, _web_search_rewrites_by_block(web_search_strings, web_search_rewrites) ) + self._write_back_output_tool_calls( + response_content, + scanned_tool_calls=tool_calls_to_check, + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + guardrail_name=guardrail_to_apply.guardrail_name, + ) verbose_proxy_logger.debug("Anthropic Messages: Processed output response: %s", response) @@ -1717,6 +1741,43 @@ class AnthropicMessagesHandler(BaseTranslation): for block_idx, rewrites in rewrites_by_block.items(): _write_back_web_search_block(_content_block_at(response_content, block_idx), rewrites) + @classmethod + def _write_back_output_tool_calls( + cls, + response_content: Sequence[object], + *, + scanned_tool_calls: Sequence[object], + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + returned_tool_calls: Sequence[object] | None, + guardrail_name: str | None, + ) -> None: + """The nth returned call belongs to the nth ``tool_use`` block, the order the calls were scanned + in. A returned list of another length cannot be placed, so the scanned calls stand, as on the + request side. Rewritten arguments that are not a JSON object cannot become a ``tool_use`` input, + so the response fails instead of carrying the arguments the guardrail meant to replace.""" + post_guardrail_tool_calls: Final = _tool_call_shapes( + returned_tool_calls + if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls) + else scanned_tool_calls + ) + tool_use_blocks: Final = tuple( + block + for block in response_content + if (fields := cls._output_block_fields(block)) is not None and fields[0] == "tool_use" + ) + if len(tool_use_blocks) != len(pre_guardrail_tool_calls): + return + rewrites: Final = tuple( + (block, after, _rewritten_tool_use_input(after.arguments)) + for block, before, after in zip(tool_use_blocks, pre_guardrail_tool_calls, post_guardrail_tool_calls) + if before != after + ) + if any(rewritten_input is None for _, _, rewritten_input in rewrites): + raise unappliable_request_rewrite(guardrail_name) + for block, after, rewritten_input in rewrites: + if rewritten_input is not None: + _write_back_tool_use_block(block, after, rewritten_input) + @staticmethod def _get_response_content(response: object) -> Sequence[object]: """Extract content list from a dict or object response.""" diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 7e052796545..7ae0b50d130 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -34,6 +34,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPIRequest, GenericGuardrailAPIResponse, GuardrailToolParam, + tool_calls_to_write_back, ) from litellm.types.utils import GenericGuardrailAPIInputs @@ -45,6 +46,9 @@ GUARDRAIL_NAME: Final = "generic_guardrail_api" _CHAT_ROW: Final[TypeAdapter[AllMessageValues]] = TypeAdapter(AllMessageValues) _SENT_ROW: Final = TypeAdapter(Mapping[str, object]) _SENT_ROWS: Final = TypeAdapter(Sequence[Mapping[str, object]]) +_SENT_TOOL_CALLS: Final[TypeAdapter[tuple[Mapping[str, object], ...] | None]] = TypeAdapter( + tuple[Mapping[str, object], ...] | None +) _JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) # Headers whose values are forwarded as-is (case-insensitive). Glob patterns supported (e.g. x-stainless-*, x-litellm*). @@ -442,12 +446,16 @@ class GenericGuardrailAPI(CustomGuardrail): tools: list[ChatCompletionToolParam] | None, structured_messages: Sequence[AllMessageValues] | None, shown_messages: Sequence[AllMessageValues] | None, + sent_tool_calls: Sequence[Mapping[str, object]] | None, guardrail_response: GenericGuardrailAPIResponse, ) -> GenericGuardrailAPIInputs: # Action is NONE or no modifications needed return_inputs: Final = GenericGuardrailAPIInputs(texts=texts) if guardrail_response.texts: return_inputs["texts"] = guardrail_response.texts + rewritten_tool_calls: Final = tool_calls_to_write_back(sent_tool_calls, guardrail_response.tool_calls) + if rewritten_tool_calls is not None: + return_inputs["tool_calls"] = list(rewritten_tool_calls) if guardrail_response.images: return_inputs["images"] = guardrail_response.images elif images: @@ -601,11 +609,8 @@ class GenericGuardrailAPI(CustomGuardrail): # Make the API request # Use mode="json" to ensure all iterables are converted to lists - response: Final = await self._post_check( - self._request_body(guardrail_request, structured_messages), - headers, - self.num_retries, - ) + guardrail_request_body: Final = self._request_body(guardrail_request, structured_messages) + response: Final = await self._post_check(guardrail_request_body, headers, self.num_retries) response.raise_for_status() response_json: Final = response.json() @@ -632,6 +637,7 @@ class GenericGuardrailAPI(CustomGuardrail): tools=tools, structured_messages=structured_messages, shown_messages=guardrail_request.structured_messages, + sent_tool_calls=_SENT_TOOL_CALLS.validate_python(guardrail_request_body.get("tool_calls")), guardrail_response=guardrail_response, ) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index 4e4983c53f0..2e1b366cef8 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -1,8 +1,8 @@ from collections.abc import Mapping, Sequence from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles -from pydantic import ConfigDict, Field, SerializerFunctionWrapHandler, model_serializer -from typing_extensions import TypedDict +from pydantic import ConfigDict, Field, SerializerFunctionWrapHandler, TypeAdapter, model_serializer, with_config +from typing_extensions import ReadOnly, TypedDict from litellm.types.guardrails import GuardrailNumRetries from litellm.types.llms.base import LiteLLMBaseModel @@ -188,12 +188,65 @@ def structured_messages_from_response(value: object) -> Sequence[AllMessageValue return cast("Sequence[AllMessageValues]", value) # cast-ok: JSON rows checked for a role, the same trust texts get +@with_config(ConfigDict(extra="allow")) +class GuardrailToolCallFunction(TypedDict): + name: ReadOnly[str] + arguments: ReadOnly[str] + + +@with_config(ConfigDict(extra="allow")) +class GuardrailToolCall(TypedDict): + """A rewritten tool call, one per call the guardrail received. The call's own ``id`` and + ``type`` are kept by the proxy, so ``function`` is all a guardrail has to send back.""" + + function: ReadOnly[GuardrailToolCallFunction] + + +_TOOL_CALLS: Final = TypeAdapter(tuple[GuardrailToolCall, ...]) +_SENT_FUNCTION: Final = TypeAdapter(Mapping[str, object]) +_RESPONSE_FIELDS: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) + + +def tool_calls_from_response(value: object) -> tuple[GuardrailToolCall, ...] | None: + """A malformed ``tool_calls`` raises, so the guardrail's ``fail_on_error`` setting decides the + request's fate instead of the unrewritten calls going through.""" + return None if value is None else _TOOL_CALLS.validate_python(value) + + +def _with_returned_function(sent_call: Mapping[str, object], returned_call: GuardrailToolCall) -> Mapping[str, object]: + sent_function: Final = _SENT_FUNCTION.validate_python(sent_call.get("function") or {}) + returned_function: Final = returned_call["function"] + return { + **sent_call, + "function": {**sent_function, "name": returned_function["name"], "arguments": returned_function["arguments"]}, + } + + +def tool_calls_to_write_back( + sent_calls: Sequence[Mapping[str, object]] | None, returned_calls: Sequence[GuardrailToolCall] | None +) -> tuple[ChatCompletionToolCallChunk, ...] | None: + """The calls as the guardrail received them, carrying the ``function`` it handed back, so every + endpoint keeps each call's own ``id`` and ``type``. The guardrail answers one call per call received, + in the same order, and an empty list is no rewrite. Any other length cannot be matched to the calls + by position, which is a guardrail error rather than a reason to send the calls unrewritten.""" + if not returned_calls: + return None + sent_count: Final = 0 if sent_calls is None else len(sent_calls) + if sent_calls is None or len(returned_calls) != sent_count: + raise ValueError(f"the guardrail returned {len(returned_calls)} tool_calls for the {sent_count} it received") + return cast( # cast-ok: the rows the proxy sent as tool calls, each with a pydantic-validated function + tuple[ChatCompletionToolCallChunk, ...], + tuple(_with_returned_function(sent, returned) for sent, returned in zip(sent_calls, returned_calls)), + ) + + class GenericGuardrailAPIResponse: """Response model for the Generic Guardrail API""" texts: list[str] | None images: list[str] | None tools: list[GuardrailToolParam] | None + tool_calls: Sequence[GuardrailToolCall] | None structured_messages: Sequence[AllMessageValues] | None action: str blocked_reason: str | None @@ -208,12 +261,14 @@ class GenericGuardrailAPIResponse: tools: list[GuardrailToolParam] | None = None, stream_holdback_chars: list[int] | None = None, structured_messages: Sequence[AllMessageValues] | None = None, + tool_calls: Sequence[GuardrailToolCall] | None = None, ) -> None: self.action = action self.blocked_reason = blocked_reason self.texts = texts self.images = images self.tools = tools + self.tool_calls = tool_calls self.structured_messages = structured_messages # Number of trailing chars, indexed the same as ``texts``, that the # framework must withhold from streaming emission until the next @@ -226,6 +281,7 @@ class GenericGuardrailAPIResponse: stream_holdback_chars: Final = ( [coerce_stream_holdback_value(value) for value in raw_holdback] if isinstance(raw_holdback, list) else None ) + fields: Final = _RESPONSE_FIELDS.validate_python(data) return cls( action=data.get("action", "NONE"), blocked_reason=data.get("blocked_reason"), @@ -234,4 +290,5 @@ class GenericGuardrailAPIResponse: tools=data.get("tools"), stream_holdback_chars=stream_holdback_chars, structured_messages=structured_messages_from_response(data.get("structured_messages")), + tool_calls=tool_calls_from_response(fields.get("tool_calls")), ) diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 8affe2e136e..e651c08d986 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -8,7 +8,8 @@ with guardrail transformations, specifically testing edge cases with empty choic import json import logging import re -from typing import Any, Literal, Optional, Final +from collections.abc import Callable, Sequence +from typing import Any, Literal, Optional, Final, cast from unittest.mock import MagicMock, patch import pytest @@ -21,6 +22,8 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, ) from litellm.llms.anthropic.messages.adapters.transformation import truncate_tool_name +from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite +from litellm.types.llms.anthropic import AnthropicResponseContentBlockText, AnthropicResponseContentBlockToolUse from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs @@ -274,12 +277,42 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: @staticmethod def _ended_sse_chunks() -> list: events = [ - ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), - ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), - ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello "}}), - ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world"}}), + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello "}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world"}}, + ), ("content_block_stop", {"type": "content_block_stop", "index": 0}), - ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + ), ("message_stop", {"type": "message_stop"}), ] return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] @@ -376,13 +409,58 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: @staticmethod def _ended_tool_use_sse_chunks() -> list: events = [ - ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), - ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}), - ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}), - ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}), - ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}), + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}, + }, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}, + }, + ), ("content_block_stop", {"type": "content_block_stop", "index": 0}), - ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + ), ("message_stop", {"type": "message_stop"}), ] return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] @@ -403,7 +481,8 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: json.loads(line[len("data:") :].strip())["delta"]["partial_json"] for chunk in chunks for line in chunk.decode().split("\n") - if line.startswith("data:") and json.loads(line[len("data:") :].strip()).get("type") == "content_block_delta" + if line.startswith("data:") + and json.loads(line[len("data:") :].strip()).get("type") == "content_block_delta" ] @pytest.mark.asyncio @@ -496,7 +575,10 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: logging_obj: Optional[LiteLLMLoggingObj] = None, ) -> GenericGuardrailAPIInputs: for tool_call in inputs.get("tool_calls", []): - if isinstance(tool_call, ChatCompletionMessageToolCall) and "persimmon" in tool_call.function.arguments: + if ( + isinstance(tool_call, ChatCompletionMessageToolCall) + and "persimmon" in tool_call.function.arguments + ): tool_call.function.arguments = '{"fruit": "[MASKED]"}' return inputs @@ -2602,9 +2684,7 @@ class TestAnthropicMessagesScanOnlyToolResults: {"role": "user", "content": "weather in Paris?"}, { "role": "assistant", - "content": [ - {"type": "tool_use", "id": "tu1", "name": "get_weather", "input": {"city": "Paris"}} - ], + "content": [{"type": "tool_use", "id": "tu1", "name": "get_weather", "input": {"city": "Paris"}}], }, { "role": "user", @@ -2949,9 +3029,7 @@ class TestStructuredWriteBackKeepsToolResults: ], ids=["tool_reference", "empty"], ) - async def test_tool_result_stays_right_after_its_tool_use( - self, tool_result_content, expected_written_back_content - ): + async def test_tool_result_stays_right_after_its_tool_use(self, tool_result_content, expected_written_back_content): handler = AnthropicMessagesHandler() data = {"model": "claude-fable-5", "messages": self._claude_code_tool_search_turns(tool_result_content)} @@ -3458,3 +3536,172 @@ def test_reading_a_dropped_budget_for_a_guardrail_logs_no_drop_warning(caplog: p assert messages == [{"role": "user", "content": "hi"}] assert not [record for record in caplog.records if "Dropping" in record.getMessage()] + + +def _tool_call_returning_guardrail(returned: Callable[[Sequence[object]], object]) -> CustomGuardrail: + class ReturnToolCalls(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + return cast( + GenericGuardrailAPIInputs, + { + **inputs, + "texts": [f"[SCANNED] {text}" for text in inputs["texts"]], + "tool_calls": returned(inputs["tool_calls"]), + }, + ) + + return ReturnToolCalls(guardrail_name="tool-call-rewriter") + + +def _with_arguments(tool_call: object, arguments: str) -> dict[str, object]: + fields: Final = cast(dict[str, object], tool_call) + function: Final = cast(dict[str, object], fields["function"]) + return {**fields, "function": {**function, "arguments": arguments}} + + +def _redacting(tool_calls: Sequence[object]) -> list[dict[str, object]]: + return [_with_arguments(tool_call, '{"cmd": "echo [REDACTED]"}') for tool_call in tool_calls] + + +def _dict_tool_use_response() -> dict[str, object]: + return { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [ + {"type": "text", "text": "Running it"}, + {"type": "tool_use", "id": "toolu_1", "name": "run_shell", "input": {"cmd": "echo SECRET-123"}}, + ], + "stop_reason": "tool_use", + } + + +def _object_tool_use_response() -> dict[str, object]: + return { + **_dict_tool_use_response(), + "content": [ + AnthropicResponseContentBlockText(type="text", text="Running it"), + AnthropicResponseContentBlockToolUse( + type="tool_use", id="toolu_1", name="run_shell", input={"cmd": "echo SECRET-123"} + ), + ], + } + + +def _tool_use_block(response: dict[str, object], index: int = 1) -> dict[str, object]: + block: Final = cast(list[object], response["content"])[index] + return cast( + dict[str, object], block.model_dump() if isinstance(block, AnthropicResponseContentBlockToolUse) else block + ) + + +class TestOutputToolCallRewrites: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "response", [_dict_tool_use_response(), _object_tool_use_response()], ids=["dict_blocks", "object_blocks"] + ) + async def test_rewritten_arguments_land_in_the_tool_use_input(self, response: dict[str, object]) -> None: + await AnthropicMessagesHandler().process_output_response( + response=cast(Any, response), + guardrail_to_apply=_tool_call_returning_guardrail(_redacting), + litellm_logging_obj=MagicMock(), + ) + + block: Final = _tool_use_block(response) + assert (block["id"], block["name"], block["input"]) == ("toolu_1", "run_shell", {"cmd": "echo [REDACTED]"}) + assert "SECRET-123" not in json.dumps(response, default=lambda block: block.model_dump()) + + @pytest.mark.asyncio + async def test_a_rewritten_name_lands_on_the_tool_use_block(self) -> None: + response: Final = _dict_tool_use_response() + + def rename(tool_calls: Sequence[object]) -> list[dict[str, object]]: + fields: Final = cast(dict[str, object], tool_calls[0]) + function: Final = cast(dict[str, object], fields["function"]) + return [{**fields, "function": {**function, "name": "run_safely"}}] + + await AnthropicMessagesHandler().process_output_response( + response=cast(Any, response), + guardrail_to_apply=_tool_call_returning_guardrail(rename), + litellm_logging_obj=MagicMock(), + ) + + assert _tool_use_block(response)["name"] == "run_safely" + assert _tool_use_block(response)["input"] == {"cmd": "echo SECRET-123"} + + @pytest.mark.asyncio + async def test_each_returned_call_lands_on_the_tool_use_block_in_the_same_position(self) -> None: + response: Final = _dict_tool_use_response() + content: Final = cast(list[object], response["content"]) + content.append({"type": "text", "text": "and then"}) + content.append({"type": "tool_use", "id": "toolu_2", "name": "lookup", "input": {"query": "SECRET-123"}}) + + def redact_second_only(tool_calls: Sequence[object]) -> list[object]: + return [tool_calls[0], _with_arguments(tool_calls[1], '{"query": "[REDACTED]"}')] + + await AnthropicMessagesHandler().process_output_response( + response=cast(Any, response), + guardrail_to_apply=_tool_call_returning_guardrail(redact_second_only), + litellm_logging_obj=MagicMock(), + ) + + assert _tool_use_block(response, 1)["input"] == {"cmd": "echo SECRET-123"} + assert _tool_use_block(response, 3)["input"] == {"query": "[REDACTED]"} + assert [block["text"] for block in content if cast(dict[str, object], block)["type"] == "text"] == [ + "[SCANNED] Running it", + "[SCANNED] and then", + ] + + @pytest.mark.asyncio + async def test_rewritten_arguments_that_are_not_a_json_object_fail_the_response(self) -> None: + response: Final = _dict_tool_use_response() + + def break_arguments(tool_calls: Sequence[object]) -> list[dict[str, object]]: + return [_with_arguments(tool_calls[0], "echo [REDACTED]")] + + with pytest.raises(UnappliableRequestRewrite, match="tool-call-rewriter"): + await AnthropicMessagesHandler().process_output_response( + response=cast(Any, response), + guardrail_to_apply=_tool_call_returning_guardrail(break_arguments), + litellm_logging_obj=MagicMock(), + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "returned", + [lambda tool_calls: [], lambda tool_calls: list(_redacting(tool_calls)) * 2, lambda tool_calls: None], + ids=["empty", "doubled", "none"], + ) + async def test_a_returned_list_that_cannot_be_placed_leaves_the_tool_use_input_as_scanned( + self, returned: Callable[[Sequence[object]], object] + ) -> None: + response: Final = _dict_tool_use_response() + + await AnthropicMessagesHandler().process_output_response( + response=cast(Any, response), + guardrail_to_apply=_tool_call_returning_guardrail(returned), + litellm_logging_obj=MagicMock(), + ) + + assert _tool_use_block(response)["input"] == {"cmd": "echo SECRET-123"} + assert _tool_use_block(response, 0) == {"type": "text", "text": "[SCANNED] Running it"} + + @pytest.mark.asyncio + async def test_echoed_tool_calls_leave_the_tool_use_block_untouched(self) -> None: + response: Final = _dict_tool_use_response() + original_input: Final = _tool_use_block(response)["input"] + + await AnthropicMessagesHandler().process_output_response( + response=cast(Any, response), + guardrail_to_apply=_tool_call_returning_guardrail(list), + litellm_logging_obj=MagicMock(), + ) + + assert _tool_use_block(response)["input"] is original_input diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index b5f794cb136..995a98612c6 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -2692,3 +2692,111 @@ async def test_during_mcp_call_blocked_raises() -> None: ) assert exc_info.value.message == "credential detected in tool input" + + +def _run_shell_call(arguments: str) -> dict[str, object]: + return {"id": "call_1", "type": "function", "function": {"name": "run_shell", "arguments": arguments}} + + +def _guardrail_answering(body: dict[str, object], fail_on_error: bool = True) -> GenericGuardrailAPI: + def answer(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=body, request=request) + + return GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="tool-call-rewriter", + event_hook="post_call", + default_on=True, + fail_on_error=fail_on_error, + async_handler=AsyncHTTPHandler(transport=httpx.MockTransport(answer)), + ) + + +async def _tool_calls_after(guardrail: GenericGuardrailAPI) -> GenericGuardrailAPIInputs: + return await guardrail.apply_guardrail( + inputs={"texts": [], "tool_calls": [_run_shell_call('{"cmd": "echo SECRET-123"}')]}, + request_data={}, + input_type="response", + ) + + +class TestToolCallsInResponse: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "returned_call", + [ + _run_shell_call('{"cmd": "echo [REDACTED]"}'), + {"function": {"name": "run_shell", "arguments": '{"cmd": "echo [REDACTED]"}'}}, + { + "id": "ignored", + "type": "ignored", + "function": {"name": "run_shell", "arguments": '{"cmd": "echo [REDACTED]"}'}, + }, + ], + ids=["as_received", "function_only", "foreign_id_and_type"], + ) + async def test_rewritten_tool_calls_reach_the_caller_with_the_calls_own_id_and_type( + self, returned_call: dict[str, object] + ) -> None: + guardrail: Final = _guardrail_answering({"action": "GUARDRAIL_INTERVENED", "tool_calls": [returned_call]}) + + guardrailed: Final = await _tool_calls_after(guardrail) + + assert guardrailed["tool_calls"] == [_run_shell_call('{"cmd": "echo [REDACTED]"}')] + + @pytest.mark.asyncio + async def test_a_rewritten_name_reaches_the_caller(self) -> None: + guardrail: Final = _guardrail_answering( + {"action": "GUARDRAIL_INTERVENED", "tool_calls": [{"function": {"name": "run_safely", "arguments": "{}"}}]} + ) + + guardrailed: Final = await _tool_calls_after(guardrail) + + assert guardrailed["tool_calls"] == [ + {"id": "call_1", "type": "function", "function": {"name": "run_safely", "arguments": "{}"}} + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "body", + [ + {"action": "NONE"}, + {"action": "NONE", "tool_calls": None}, + {"action": "GUARDRAIL_INTERVENED", "texts": ["[MASKED]"], "tool_calls": []}, + ], + ids=["absent", "null", "empty"], + ) + async def test_an_answer_without_tool_calls_leaves_the_key_out(self, body: dict[str, object]) -> None: + guardrailed: Final = await _tool_calls_after(_guardrail_answering(body)) + + assert "tool_calls" not in guardrailed + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "tool_calls", + [ + "not a list", + ["not a tool call"], + [{"id": "call_1", "type": "function", "function": None}], + [{"function": {"name": "run_shell"}}], + [{"function": {"name": "run_shell", "arguments": {"cmd": "echo [REDACTED]"}}}], + [_run_shell_call("{}"), _run_shell_call("{}")], + ], + ids=["not_a_list", "not_a_call", "null_function", "no_arguments", "object_arguments", "one_call_too_many"], + ) + async def test_a_tool_calls_answer_that_cannot_be_applied_fails_the_guardrail(self, tool_calls: object) -> None: + guardrail: Final = _guardrail_answering({"action": "GUARDRAIL_INTERVENED", "tool_calls": tool_calls}) + + with pytest.raises(Exception, match="Generic Guardrail API failed"): + await _tool_calls_after(guardrail) + + @pytest.mark.asyncio + async def test_fail_on_error_off_sends_the_calls_as_received_when_the_answer_cannot_be_applied(self) -> None: + guardrail: Final = _guardrail_answering( + {"action": "GUARDRAIL_INTERVENED", "tool_calls": [_run_shell_call("{}"), _run_shell_call("{}")]}, + fail_on_error=False, + ) + + guardrailed: Final = await _tool_calls_after(guardrail) + + assert guardrailed["tool_calls"] == [_run_shell_call('{"cmd": "echo SECRET-123"}')]