From 5fa693896fea7ac6ab2455bbfbfbb5f11171eef7 Mon Sep 17 00:00:00 2001 From: Chetan Soni Date: Fri, 11 Sep 2026 01:21:15 -0700 Subject: [PATCH 1/2] fix(guardrails): stop 500s on POST /responses when a guardrail rewrites input --- .../guardrail_translation/handler.py | 13 +- .../crowdstrike_aidr/crowdstrike_aidr.py | 8 +- .../proxy/policy_engine/pipeline_executor.py | 9 ++ .../guardrail_hooks/test_crowdstrike_aidr.py | 117 ++++++++++++++++++ 4 files changed, 143 insertions(+), 4 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 33e0a0a923f..2fe11d9f7bd 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -497,9 +497,14 @@ class OpenAIResponsesHandler(BaseTranslation): guardrailed_texts: Final = guardrailed_inputs.get("texts") or () data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param else: + rewritten_texts: Final = guardrailed_inputs.get("texts") or () + if len(rewritten_texts) != len(extracted.task_mappings): + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown") await self._apply_guardrail_responses_to_input( messages=input_data, - responses=guardrailed_inputs.get("texts") or (), + responses=rewritten_texts, task_mappings=extracted.task_mappings, ) verbose_proxy_logger.debug("OpenAI Responses API: Processed input messages: %s", data.get("input")) @@ -635,10 +640,12 @@ class OpenAIResponsesHandler(BaseTranslation): """ Apply guardrail responses back to input messages. + ``responses`` pairs positionally with ``task_mappings``; the caller rejects + the request when the two disagree, so this never has to guess an alignment. + Override this method to customize how responses are applied. """ - for task_idx, guardrail_response in enumerate(responses): - mapping = task_mappings[task_idx] + for guardrail_response, mapping in zip(responses, task_mappings): msg_idx = cast(int, mapping[0]) content_idx_optional = cast(int | None, mapping[1]) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index f7f500b1adc..8fed1f906e5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -18,6 +18,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_tool_message_for_guardrail, ) from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, get_async_httpx_client, httpxSpecialProvider, ) @@ -261,6 +262,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): fail_on_error: bool | None = True, streaming_end_of_stream_only: bool | None = None, streaming_sampling_rate: int | None = None, + async_handler: AsyncHTTPHandler | None = None, **kwargs, ) -> None: """ @@ -273,9 +275,13 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): streaming_end_of_stream_only (bool | None): Scan streamed output once at end of stream instead of every streaming_sampling_rate chunks. Defaults to False. streaming_sampling_rate (int | None): Scan the accumulated streamed output every Nth chunk. Defaults to 5. + async_handler (AsyncHTTPHandler | None): HTTP client to call AI Guard with. Defaults to the shared + guardrail-callback client. **kwargs: Additional arguments passed to the CustomGuardrail base class. """ - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler = async_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) self.fail_on_error = True if fail_on_error is None else fail_on_error self._set_streaming_params( CrowdStrikeAIDRGuardrailConfigModelOptionalParams( diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index ed193c7f434..ad45781d5d2 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -58,6 +58,15 @@ class UndeliverableStreamRewrite(Exception): self.guardrail_name: Final = guardrail_name +class UnappliableRequestRewrite(Exception): + def __init__(self, guardrail_name: str) -> None: + super().__init__( + f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, " + "so the request was rejected rather than sent unrewritten" + ) + self.guardrail_name: Final = guardrail_name + + def _tool_call_shape(tool_call: object) -> tuple[object, object]: plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call function: Final = plain.get("function") if isinstance(plain, Mapping) else None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index a07157396df..a80dae27d01 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -1,3 +1,6 @@ +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Final, cast from unittest.mock import patch import httpx @@ -7,6 +10,9 @@ from pydantic import ValidationError import litellm from litellm.exceptions import Timeout +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.crowdstrike_aidr.crowdstrike_aidr import ( @@ -1719,3 +1725,114 @@ async def test_streaming_params_from_config_control_output_scan_cadence( handler = _initialize_from_config(mode="post_call", **configured) assert await _guard_calls_for_stream(handler, list("ABCDEFGHIJ")) == expected_calls + + +@asynccontextmanager +async def _guardrail_transforming_into(messages: list[dict[str, object]]) -> AsyncIterator[CrowdStrikeAIDRHandler]: + """A guardrail whose AI Guard returns ``messages`` as its rewrite.""" + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "result": { + "blocked": False, + "transformed": True, + "guard_output": {"messages": messages}, + }, + }, + request=request, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + handler: Final = AsyncHTTPHandler() + handler.client = client + yield CrowdStrikeAIDRHandler( + mode="pre_call", + guardrail_name="crowdstrike-aidr-guard", + api_key="pts_crowdstrike_tokenid", + api_base="https://api.crowdstrike.com/aidr/aiguard", + async_handler=handler, + ) + + +class _MessageShapedGuardrail(CustomGuardrail): + """Returns one text per chat message and no ``structured_messages`` rewrite. + + Prompt Security and friends scan messages rather than Responses text parts, + which is the shape that outnumbers the endpoint's own bookkeeping. + """ + + def __init__(self, redacted: str) -> None: + super().__init__(guardrail_name="message-shaped") + self.redacted: Final = redacted + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: str, + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + messages: Final = inputs.get("structured_messages") or () + return {"texts": [self.redacted for _ in messages]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("case", "instructions", "responses_input"), + [ + ( + "instructions add a system message", + "be terse", + [{"role": "user", "content": [{"type": "input_text", "text": "my ssn is 078-05-1120"}]}], + ), + ( + "tool items add messages that carry no text", + None, + [ + {"role": "user", "content": [{"type": "input_text", "text": "my ssn is 078-05-1120"}]}, + {"type": "function_call", "call_id": "c1", "name": "get_x", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "c1", "output": "42"}, + ], + ), + ], +) +async def test_unalignable_rewrite_is_rejected_never_sent_unredacted( + case: str, + instructions: str | None, + responses_input: list[dict[str, object]], +) -> None: + """An unalignable rewrite must fail the request, not forward the raw prompt. + + Skipping the write-back would hand the model the unredacted text, so a + guardrail could be bypassed by adding ``instructions`` or a tool call. + """ + from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite + + data: dict[str, object] = {"model": "gpt-4o", "input": responses_input} + if instructions is not None: + data["instructions"] = instructions + + with pytest.raises(UnappliableRequestRewrite): + await OpenAIResponsesHandler().process_input_messages( + data=data, + guardrail_to_apply=_MessageShapedGuardrail("my ssn is "), + ) + + assert "078-05-1120" in str(responses_input), case + + +@pytest.mark.asyncio +async def test_aligned_rewrite_is_written_back() -> None: + """Matching counts must still redact the input in place.""" + responses_input: list[dict[str, object]] = [ + {"role": "user", "content": [{"type": "input_text", "text": "my ssn is 078-05-1120"}]} + ] + + await OpenAIResponsesHandler().process_input_messages( + data={"model": "gpt-4o", "input": responses_input}, + guardrail_to_apply=_MessageShapedGuardrail("my ssn is "), + ) + + assert cast(list, responses_input[0]["content"])[0]["text"] == "my ssn is " From e4030597d80ad6e2b8d1a946a9dbbb1082ed6285 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:47:18 -0700 Subject: [PATCH 2/2] test(guardrails): prove the structured write-back lands CrowdStrike redactions on Responses instructions and tool items --- .../guardrail_hooks/test_crowdstrike_aidr.py | 58 ++++++++++++++++++- 1 file changed, 55 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index a80dae27d01..a9ca13a463d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -1,6 +1,7 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from typing import Final, cast +import json from unittest.mock import patch import httpx @@ -1728,17 +1729,28 @@ async def test_streaming_params_from_config_control_output_scan_cadence( @asynccontextmanager -async def _guardrail_transforming_into(messages: list[dict[str, object]]) -> AsyncIterator[CrowdStrikeAIDRHandler]: - """A guardrail whose AI Guard returns ``messages`` as its rewrite.""" +async def _guardrail_redacting(secret: str, replacement: str) -> AsyncIterator[CrowdStrikeAIDRHandler]: + def redacted(content: object) -> object: + if isinstance(content, str): + return content.replace(secret, replacement) + if isinstance(content, list): + return [ + {**part, "text": redacted(part["text"])} if isinstance(part, dict) and "text" in part else part + for part in content + ] + return content def respond(request: httpx.Request) -> httpx.Response: + sent: Final = json.loads(request.content)["guard_input"]["messages"] return httpx.Response( status_code=200, json={ "result": { "blocked": False, "transformed": True, - "guard_output": {"messages": messages}, + "guard_output": { + "messages": [{**message, "content": redacted(message.get("content"))} for message in sent] + }, }, }, request=request, @@ -1836,3 +1848,43 @@ async def test_aligned_rewrite_is_written_back() -> None: ) assert cast(list, responses_input[0]["content"])[0]["text"] == "my ssn is " + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("case", "responses_input", "redacted_input"), + [ + ( + "instructions add a system message", + [{"role": "user", "content": [{"type": "input_text", "text": "my ssn is 078-05-1120"}]}], + [{"role": "user", "content": [{"type": "input_text", "text": "my ssn is "}]}], + ), + ( + "tool items sit between two user turns", + [ + {"role": "user", "content": [{"type": "input_text", "text": "hello"}]}, + {"type": "function_call", "call_id": "c1", "name": "get_x", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "c1", "output": "42"}, + {"role": "user", "content": [{"type": "input_text", "text": "my ssn is 078-05-1120"}]}, + ], + [ + {"role": "user", "content": [{"type": "input_text", "text": "hello"}]}, + {"type": "function_call", "call_id": "c1", "name": "get_x", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "c1", "output": "42"}, + {"role": "user", "content": [{"type": "input_text", "text": "my ssn is "}]}, + ], + ), + ], +) +async def test_structured_rewrite_lands_on_shapes_the_flat_path_cannot_align( + case: str, + responses_input: list[dict[str, object]], + redacted_input: list[dict[str, object]], +) -> None: + data: dict[str, object] = {"model": "gpt-5.6", "instructions": "be terse", "input": responses_input} + + async with _guardrail_redacting("078-05-1120", "") as guardrail: + await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail) + + assert data["input"] == redacted_input, case + assert data["instructions"] == "be terse", case