From a248ff2c6ea712b0f8cb9f03bc327701c8dd63ab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:14:50 -0700 Subject: [PATCH] fix(guardrails): fail open on Responses tool-call rewrites in another shape and keep item names Guardrails that hand back tool_calls in their own shape (vendor JSON, user code output) raised a KeyError on the non-stream Responses write-back. Returned tool calls are now validated before comparison; a shape or count that does not line up leaves every tool-call item unchanged and logs a warning naming the guardrail. A tool call's name is written back only when the guardrail changed it, so a nameless custom_tool_call no longer picks up the custom_tool placeholder. --- .../guardrail_translation/handler.py | 85 ++++++++++++++---- ...test_openai_responses_guardrail_handler.py | 90 +++++++++++++++++++ 2 files changed, 159 insertions(+), 16 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index fd0d7b66bee..8c8294ad2fc 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -37,7 +37,7 @@ from itertools import accumulate, chain, repeat from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Union, cast -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -105,6 +105,19 @@ class _ToolCallShape(NamedTuple): arguments: str +class _ToolCallFunctionFields(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str | None = None + arguments: str = "" + + +class _ToolCallFields(BaseModel): + model_config = ConfigDict(frozen=True) + + function: _ToolCallFunctionFields + + def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]: return tuple( _ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", "")) @@ -112,6 +125,47 @@ def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tupl ) +def _returned_tool_call_shape(tool_call: object) -> _ToolCallShape | None: + payload: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call + try: + fields: Final = _ToolCallFields.model_validate(payload) + except ValidationError: + return None + return _ToolCallShape(name=fields.function.name, arguments=fields.function.arguments) + + +def _post_guardrail_tool_call_shapes( + returned_tool_calls: Sequence[object] | None, + pre_guardrail_tool_calls: tuple[_ToolCallShape, ...], + guardrail_name: str | None, +) -> tuple[_ToolCallShape, ...]: + if not pre_guardrail_tool_calls: + return pre_guardrail_tool_calls + if returned_tool_calls is None or len(returned_tool_calls) != len(pre_guardrail_tool_calls): + verbose_proxy_logger.warning( + "OpenAI Responses API: guardrail %s returned %s tool calls for the %d scanned, " + "leaving the tool call output items unchanged", + guardrail_name, + "no" if returned_tool_calls is None else len(returned_tool_calls), + len(pre_guardrail_tool_calls), + ) + return pre_guardrail_tool_calls + returned_shapes: Final = tuple(_returned_tool_call_shape(tool_call) for tool_call in returned_tool_calls) + validated_shapes: Final = tuple(shape for shape in returned_shapes if shape is not None) + if len(validated_shapes) != len(returned_shapes): + verbose_proxy_logger.warning( + "OpenAI Responses API: guardrail %s returned tool calls without a function name and arguments, " + "leaving the tool call output items unchanged", + guardrail_name, + ) + return pre_guardrail_tool_calls + return validated_shapes + + +def _tool_call_rewrite(before: _ToolCallShape, after: _ToolCallShape) -> _ToolCallShape: + return _ToolCallShape(name=after.name if after.name != before.name else None, arguments=after.arguments) + + class ResponseOutputEnvelope(TypedDict, total=False): """Dict form of a Responses API response, as far as guardrail write-back reads it.""" @@ -698,7 +752,11 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") + post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes( + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name, + ) # Step 3: Map guardrail responses back to original response structure await self._apply_guardrail_responses_to_output( @@ -709,11 +767,7 @@ class OpenAIResponsesHandler(BaseTranslation): self._write_tool_call_rewrites_to_output( tool_call_items=tuple(item for item in response_output if _is_tool_call_output_item(item)), pre_guardrail_tool_calls=pre_guardrail_tool_calls, - post_guardrail_tool_calls=_tool_call_shapes( - returned_tool_calls - if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) - else tool_calls_to_check - ), + post_guardrail_tool_calls=post_guardrail_tool_calls, ) verbose_proxy_logger.debug("OpenAI Responses API: Processed output response: %s", response) @@ -811,11 +865,10 @@ class OpenAIResponsesHandler(BaseTranslation): ) guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") - post_guardrail_tool_calls: Final = _tool_call_shapes( - returned_tool_calls - if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check) - else tool_calls_to_check + post_guardrail_tool_calls: Final = _post_guardrail_tool_call_shapes( + returned_tool_calls=guardrailed_inputs.get("tool_calls"), + pre_guardrail_tool_calls=pre_guardrail_tool_calls, + guardrail_name=guardrail_to_apply.guardrail_name, ) # Write guardrailed texts back into the output items in-place. @@ -991,7 +1044,7 @@ class OpenAIResponsesHandler(BaseTranslation): ) rewrites_by_call_id: Final = MappingProxyType( { - call_id: after + call_id: _tool_call_rewrite(before, after) for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls) if after != before } @@ -1050,12 +1103,12 @@ class OpenAIResponsesHandler(BaseTranslation): ) -> None: if len(tool_call_items) != len(post_guardrail_tool_calls): return - for output_item, after in ( - (output_item, after) + for output_item, rewrite in ( + (output_item, _tool_call_rewrite(before, after)) for output_item, before, after in zip(tool_call_items, pre_guardrail_tool_calls, post_guardrail_tool_calls) if after != before ): - self._write_tool_call_item(output_item, after.name, after.arguments) + self._write_tool_call_item(output_item, rewrite.name, rewrite.arguments) @staticmethod def _tool_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]: diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index cf4895f637f..a4f0a77a9b6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -10,6 +10,8 @@ from collections.abc import Callable from typing import Any, List, Literal, Optional, Tuple from unittest.mock import AsyncMock, MagicMock +import logging + import pytest @@ -85,6 +87,29 @@ class PersimmonMaskingGuardrail(CustomGuardrail): return {**inputs, "tool_calls": tool_calls} +class FlatShapeGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + flat_tool_calls = [{"name": "exec", "input": "rm -rf /"} for _ in inputs.get("tool_calls", [])] + return {**inputs, "tool_calls": flat_tool_calls} + + +class DroppingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[LiteLLMLoggingObj] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "tool_calls": []} + + CUSTOM_TOOL_CALL_ITEM = { "type": "custom_tool_call", "id": "ctc_1", @@ -737,6 +762,52 @@ class TestOpenAIResponsesHandlerToolCallExtraction: assert (custom_item.name if typed else custom_item["name"]) == "exec" assert (output[0].content[0].text if typed else output[0]["content"][0]["text"]) == "running persimmon" + @staticmethod + def _custom_tool_call_response(item: dict) -> dict: + return { + "id": "resp_1", + "created_at": 1, + "model": "gpt-5.6", + "object": "response", + "status": "completed", + "output": [item], + } + + @pytest.mark.asyncio + async def test_process_output_response_ignores_tool_call_rewrites_in_another_shape(self): + handler = OpenAIResponsesHandler() + response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM)) + + result = await handler.process_output_response(response, FlatShapeGuardrail(guardrail_name="flat")) + + assert result["output"][0]["input"] == "echo persimmon" + assert result["output"][0]["name"] == "exec" + + @pytest.mark.asyncio + async def test_process_output_response_warns_when_guardrail_drops_tool_calls(self, caplog): + handler = OpenAIResponsesHandler() + response = self._custom_tool_call_response(dict(CUSTOM_TOOL_CALL_ITEM)) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await handler.process_output_response(response, DroppingGuardrail(guardrail_name="dropper")) + + assert result["output"][0]["input"] == "echo persimmon" + assert any( + "dropper" in record.getMessage() and "0 tool calls for the 1 scanned" in record.getMessage() + for record in caplog.records + ) + + @pytest.mark.asyncio + async def test_process_output_response_keeps_a_nameless_custom_tool_call_nameless(self): + handler = OpenAIResponsesHandler() + nameless_item = {key: value for key, value in CUSTOM_TOOL_CALL_ITEM.items() if key != "name"} + response = self._custom_tool_call_response(nameless_item) + + result = await handler.process_output_response(response, PersimmonMaskingGuardrail(guardrail_name="mask")) + + assert result["output"][0]["input"] == "echo [MASKED]" + assert "name" not in result["output"][0] + @pytest.mark.asyncio async def test_process_output_response_with_tool_calls(self): """Test processing output response containing function tool calls""" @@ -1469,6 +1540,25 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[5]["response"]["output"][0]["name"] == "exec" assert "arguments" not in events[5]["response"]["output"][0] + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_keep_a_nameless_custom_tool_call_nameless(self): + handler = OpenAIResponsesHandler() + events = self._ended_custom_tool_call_stream_events() + items = [events[0]["item"], events[4]["item"], events[5]["response"]["output"][0]] + for item in items: + del item["name"] + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=PersimmonMaskingGuardrail(guardrail_name="mask"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert events[3]["input"] == "echo [MASKED]" + assert events[5]["response"]["output"][0]["input"] == "echo [MASKED]" + assert all("name" not in item for item in items) + @pytest.mark.asyncio async def test_deliver_ended_stream_rewrites_syncs_typed_custom_tool_call_events(self): from litellm.types.llms.openai import (