From 574010a2cec368ba2fcffbc6fd8717161b33a3b1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 22:04:26 -0700 Subject: [PATCH] fix(responses): make guardrail input provenance O(n) and guard non-list structured_messages _input_item_provenance converted every input prefix, so an n-item request paid for n+1 full conversions. It now converts each item once, glues consecutive function_call items (plus their trailing-assistant context) into units so the transform's tool_call merging is reproduced inside the unit conversion, and verifies the unit concatenation against one full conversion, bailing to the full-conversion fallback on any mismatch. Messages from multi-item units are tainted, which keeps parallel tool calls patchable exactly like the old prefix pass while unpredicted merges fall back safely. A guardrail handing back a non-list structured_messages payload (the HiddenLayer v2 evaluation dict) previously fell through the length-mismatch fallback and 500ed converting the dict's keys as messages. The write-back is now skipped for non-list payloads, restoring the previous no-write-back behavior on the Responses surface. Also refreshes the compresr texts-mirror docstring, which still claimed the Responses translation cannot round-trip structured_messages. --- .../guardrail_translation/handler.py | 75 ++++- .../guardrail_hooks/compresr/compresr.py | 7 +- ...test_openai_responses_guardrail_handler.py | 296 ++++++++++++++++++ 3 files changed, 365 insertions(+), 13 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 8c014feeded..bd58fb4cdbc 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -29,6 +29,7 @@ Output: response.output is List[GenericResponseOutputItem] where each has: """ from collections.abc import Mapping, Sequence +from itertools import accumulate from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast @@ -119,33 +120,85 @@ def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapp return {**item, field: converted_value} # mutable-ok: request input items must stay JSON-plain dicts +def _is_function_call_item(item: object) -> bool: + return isinstance(item, Mapping) and item.get("type") in ("function_call", "custom_tool_call") + + +def _last_message_role(messages: Sequence[object]) -> str | None: + if not messages: + return None + last: Final = messages[-1] + role: Final = last.get("role") if isinstance(last, Mapping) else getattr(last, "role", None) + return role if isinstance(role, str) else None + + +def _provenance_unit_bounds( + raw_input: Sequence[object], + solo_conversions: Sequence[Sequence[object]], +) -> tuple[tuple[int, int], ...]: + trailing_roles: Final = tuple( + accumulate( + (_last_message_role(messages) for messages in solo_conversions), + lambda previous, current: current if current is not None else previous, + ) + ) + start_indexes: Final = tuple( + index + for index in range(len(raw_input)) + if index == 0 or not (_is_function_call_item(raw_input[index]) and trailing_roles[index - 1] == "assistant") + ) + return tuple(zip(start_indexes, (*start_indexes[1:], len(raw_input)))) + + def _input_item_provenance( raw_input: Sequence[object], expected_messages: Sequence[object], ) -> tuple[Mapping[int, int], frozenset[int]] | None: if not all(isinstance(item, Mapping) for item in raw_input): return None - prefixes: Final = tuple( + solo_conversions: Final = tuple( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=cast("ResponseInputParam", raw_input[:count]), # cast-ok: items checked as Mappings above + input=cast("ResponseInputParam", [item]), # cast-ok: items checked as Mappings above responses_api_request=_EMPTY_RESPONSES_REQUEST, ) - for count in range(len(raw_input) + 1) + for item in raw_input ) - if tuple(prefixes[-1]) != tuple(expected_messages): + full_conversion: Final = tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", list(raw_input)), # cast-ok: items checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + ) + if full_conversion != tuple(expected_messages): return None + units: Final = _provenance_unit_bounds(raw_input, solo_conversions) + unit_messages: Final = tuple( + tuple(solo_conversions[start]) + if end - start == 1 + else tuple( + LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=cast("ResponseInputParam", list(raw_input[start:end])), # cast-ok: checked as Mappings above + responses_api_request=_EMPTY_RESPONSES_REQUEST, + ) + ) + for start, end in units + ) + if tuple(message for messages in unit_messages for message in messages) != full_conversion: + return None + boundaries: Final = tuple(accumulate((len(messages) for messages in unit_messages), initial=0)) item_for_message: Final = MappingProxyType( { - message_index: item_index - for item_index in range(len(raw_input)) - for message_index in range(len(prefixes[item_index]), len(prefixes[item_index + 1])) + message_index: start + for unit_index, (start, end) in enumerate(units) + if end - start == 1 + for message_index in range(boundaries[unit_index], boundaries[unit_index + 1]) } ) tainted: Final = frozenset( message_index - for item_index in range(len(raw_input)) - for message_index in range(len(prefixes[item_index])) - if prefixes[item_index + 1][message_index] != prefixes[item_index][message_index] + for unit_index, (start, end) in enumerate(units) + if end - start > 1 + for message_index in range(boundaries[unit_index], boundaries[unit_index + 1]) ) return item_for_message, tainted @@ -351,6 +404,8 @@ class OpenAIResponsesHandler(BaseTranslation): original_messages: Sequence[object], structured_messages: Sequence[AllMessageValues], ) -> None: + if not isinstance(structured_messages, list): + return if _patch_rewritten_rows_into_input(data, original_messages, structured_messages): return input_items, instructions = ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 5c14d03f50e..84b328a9c60 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -916,9 +916,10 @@ class CompresrGuardrail(CustomGuardrail): def _mirror_texts_channel(input_texts: object, applied: _CompressionResult) -> list[object] | None: """Compressed content mirrored into the Responses `texts` channel. - The chat/Anthropic handlers round-trip ``structured_messages``; the - Responses translation cannot rebuild its input from chat messages and - instead writes back through ``texts``. This matches by value, so a + The chat/Anthropic/Responses handlers round-trip + ``structured_messages``; translations without that round-trip write + back through ``texts``, so the compressed content is mirrored there + too. This matches by value, so a replacement is applied only when it is unambiguous: one compression per text, and every occurrence in ``texts`` accounted for by a compressed target. Anything else is left uncompressed rather than risk a wrong or 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 87449a4b1e8..86c023dd7a4 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 @@ -1450,3 +1450,299 @@ class TestStructuredMessagesWriteBack: assert result["input"] is original_input assert [_texts(item) for item in result["input"]] == [["Hello [GUARDRAILED]"], ["Again [GUARDRAILED]"]] + + +class AllToolOutputsRewriteGuardrail(CustomGuardrail): + """Guardrail that compresses every tool-result row.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + rewritten = [ + {**m, "content": COMPRESSED_MARKER} if isinstance(m, dict) and m.get("role") == "tool" else m + for m in messages + ] + return {**inputs, "structured_messages": rewritten} + + +class AssistantRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the first assistant row's content.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "assistant") + rewritten = [{**m, "content": COMPRESSED_MARKER} if i == first else m for i, m in enumerate(messages)] + return {**inputs, "structured_messages": rewritten} + + +class DictStructuredMessagesGuardrail(CustomGuardrail): + """Guardrail that hands back a raw evaluation dict instead of a message list, + the way HiddenLayer v2 does.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "structured_messages": {"evaluation": "allowed", "messages": []}} + + +def _parallel_tool_call_input() -> list: + return [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"id": "fc_2", "type": "function_call", "call_id": "call_2", "name": "read_b", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo " * 400}, + {"type": "function_call_output", "call_id": "call_2", "output": "note " * 400}, + {"role": "user", "content": "What is the codename?"}, + ] + + +class TestProvenancePatching: + """The O(n) provenance pass must keep patching rewritten rows in place for the + shapes real agent loops produce, and fall back safely everywhere else.""" + + @pytest.mark.asyncio + async def test_parallel_tool_call_outputs_both_patched(self): + handler = OpenAIResponsesHandler() + raw_input = _parallel_tool_call_input() + fc_1, fc_2 = raw_input[0], raw_input[1] + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, AllToolOutputsRewriteGuardrail()) + + assert result["input"][0] is fc_1 + assert result["input"][1] is fc_2 + assert result["input"][2] == {"type": "function_call_output", "call_id": "call_1", "output": COMPRESSED_MARKER} + assert result["input"][3] == {"type": "function_call_output", "call_id": "call_2", "output": COMPRESSED_MARKER} + assert result["input"][4] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_assistant_turn_with_tool_call_keeps_items_verbatim(self): + handler = OpenAIResponsesHandler() + assistant_item = {"role": "assistant", "content": "Let me read the memo."} + function_call_item = { + "id": "fc_9", + "type": "function_call", + "call_id": "call_9", + "name": "read_document", + "arguments": '{"path": "memo.txt"}', + } + data = { + "model": "gpt-5.6", + "input": [ + assistant_item, + function_call_item, + {"type": "function_call_output", "call_id": "call_9", "output": "memo " * 400}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, ToolOutputRewriteGuardrail()) + + assert result["input"][0] is assistant_item + assert result["input"][1] is function_call_item + assert result["input"][2] == {"type": "function_call_output", "call_id": "call_9", "output": COMPRESSED_MARKER} + + @pytest.mark.asyncio + async def test_rewrite_of_merged_tool_call_message_falls_back(self): + handler = OpenAIResponsesHandler() + raw_input = _parallel_tool_call_input() + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, AssistantRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] + assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + + @pytest.mark.asyncio + async def test_rewrite_of_lone_function_call_message_falls_back(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo memo"}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + raw_input = data["input"] + result = await handler.process_input_messages(data, AssistantRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] + assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + + def test_provenance_bails_on_non_mapping_item(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + + assert _input_item_provenance(["not a mapping"], []) is None + + def test_provenance_bails_when_expected_messages_disagree(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + + assert _input_item_provenance([{"role": "user", "content": "hi"}], [{"role": "user", "content": "bye"}]) is None + + def test_provenance_bails_on_unpredicted_merge(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + raw_input = [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"role": "assistant", "content": "Reading the memo now."}, + ] + expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=raw_input, responses_api_request={} + ) + assert len(expected) == 1 + assert _input_item_provenance(raw_input, expected) is None + + def test_provenance_maps_and_taints_parallel_tool_calls(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + raw_input = _parallel_tool_call_input() + expected = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=raw_input, responses_api_request={} + ) + provenance = _input_item_provenance(raw_input, expected) + assert provenance is not None + item_for_message, tainted = provenance + assert tainted == {0} + assert dict(item_for_message) == {1: 2, 2: 3, 3: 4} + + +class TestDictStructuredMessagesGuard: + """A guardrail handing back a non-list structured_messages payload must not + blow up the request; the write-back is skipped instead.""" + + @pytest.mark.asyncio + async def test_list_input_survives_dict_structured_messages(self): + handler = OpenAIResponsesHandler() + original_input = [{"role": "user", "content": "Hello"}] + data = {"model": "gpt-5.6", "input": original_input} + + result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail()) + + assert result["input"] is original_input + assert result["input"] == [{"role": "user", "content": "Hello"}] + + @pytest.mark.asyncio + async def test_string_input_survives_dict_structured_messages(self): + handler = OpenAIResponsesHandler() + data = {"model": "gpt-5.6", "input": "Hello there"} + + result = await handler.process_input_messages(data, DictStructuredMessagesGuardrail()) + + assert result["input"] == "Hello there" + + +class SystemRewriteGuardrail(CustomGuardrail): + """Guardrail that rewrites the system row, the way prompt-hardening guardrails do.""" + + def __init__(self, rewritten_content: Any = COMPRESSED_MARKER): + super().__init__() + self.rewritten_content = rewritten_content + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + messages = list(inputs.get("structured_messages") or []) + first = next(i for i, m in enumerate(messages) if isinstance(m, dict) and m.get("role") == "system") + rewritten = [ + {**m, "content": self.rewritten_content} if i == first else m for i, m in enumerate(messages) + ] + return {**inputs, "structured_messages": rewritten} + + +class TestPatchEdgeBranches: + @pytest.mark.asyncio + async def test_multimodal_user_item_rewritten_through_conversion(self): + handler = OpenAIResponsesHandler() + data = { + "model": "gpt-5.6", + "input": [ + {"role": "user", "content": [{"type": "input_text", "text": "memo " * 400}]}, + {"role": "user", "content": "What is the codename?"}, + ], + } + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert _texts(result["input"][0]) == [COMPRESSED_MARKER] + assert result["input"][1] == {"role": "user", "content": "What is the codename?"} + + @pytest.mark.asyncio + async def test_instructions_rewrite_lands_in_instructions_field(self): + handler = OpenAIResponsesHandler() + user_item = {"role": "user", "content": "What is the codename?"} + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [user_item], + } + + result = await handler.process_input_messages(data, SystemRewriteGuardrail()) + + assert result["instructions"] == COMPRESSED_MARKER + assert result["input"][0] is user_item + + @pytest.mark.asyncio + async def test_non_string_instructions_rewrite_falls_back(self): + handler = OpenAIResponsesHandler() + user_item = {"role": "user", "content": "What is the codename?"} + data = { + "model": "gpt-5.6", + "instructions": "Answer from the memo only.", + "input": [user_item], + } + + result = await handler.process_input_messages( + data, SystemRewriteGuardrail(rewritten_content=[{"type": "text", "text": COMPRESSED_MARKER}]) + ) + + assert result["input"][0] is not user_item + + @pytest.mark.asyncio + async def test_unpredicted_merge_falls_back_through_patch(self): + handler = OpenAIResponsesHandler() + raw_input = [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, + {"role": "assistant", "content": "Reading the memo now."}, + {"type": "function_call_output", "call_id": "call_1", "output": "memo memo"}, + {"role": "user", "content": "memo " * 400}, + ] + data = {"model": "gpt-5.6", "input": raw_input} + + result = await handler.process_input_messages(data, StructuredRewriteGuardrail()) + + assert not any(item is original for item in result["input"] for original in raw_input) + user_items = [item for item in result["input"] if item.get("role") == "user"] + assert _texts(user_items[0]) == [COMPRESSED_MARKER] + + def test_item_rewrite_field_ignores_non_string_type(self): + from litellm.llms.openai.responses.guardrail_translation.handler import _item_rewrite_field + + assert _item_rewrite_field({"type": 123, "content": "hello"}) is None