Merge pull request #40939 from BerriAI/litellm_responses_per_message_guardrail_rewrite

fix(guardrails): write per-message guardrail rewrites back onto Responses input items

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -44,6 +44,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     scoped_structured_message_indices,
     stream_item_field,
     stream_item_fingerprint,
+    unappliable_request_rewrite,
 )
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
     AnthropicPassthroughLoggingHandler,
@@ -570,6 +571,8 @@ class AnthropicMessagesHandler(BaseTranslation):
                     preserve_system_messages=has_midturn_system_message,
                 )
             else:
+                if guardrailed_texts and len(guardrailed_texts) != len(scanned):
+                    raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
                 # Step 3: Map guardrail responses back to original message structure
                 await self._apply_guardrail_responses_to_input(
                     messages=messages,

litellm/llms/base_llm/guardrail_translation/utils.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 import json
-from collections.abc import Callable, Iterator, Sequence
-from typing import Final, TypeVar
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from typing import Final, TypeVar, cast  # noqa: TID251  # a rebuilt chat row has no typed constructor across roles
 
 from pydantic import BaseModel
 
@@ -364,3 +364,67 @@ def merge_guardrailed_scoped_messages(
                 yield from appended
 
     return list(_merged())
+
+
+def _content_part_text(part: object) -> str | None:
+    if not isinstance(part, Mapping):
+        return None
+    text: Final = part.get("text")
+    return text if isinstance(text, str) else None
+
+
+def message_slot_texts(message: Mapping[str, object]) -> tuple[str, ...]:
+    content: Final = message.get("content")
+    if isinstance(content, str):
+        return (content,)
+    if isinstance(content, list):
+        return tuple(text for part in content if (text := _content_part_text(part)) is not None)
+    return ()
+
+
+def message_text_slot_count(message: AllMessageValues) -> int:
+    return len(message_slot_texts(message))
+
+
+def _part_with_text(part: object, text: str) -> object:
+    if not isinstance(part, Mapping):
+        return part
+    return {**part, "text": text}  # mutable-ok: content parts stay JSON-plain dicts
+
+
+def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]:
+    remaining_texts: Final = iter(texts)
+    return [  # mutable-ok: message content stays a JSON list
+        _part_with_text(part, next(remaining_texts)) if _content_part_text(part) is not None else part
+        for part in content
+    ]
+
+
+def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None:
+    """Swap one rewritten text into each text slot of a chat row, in order.
+
+    A slot is a string ``content`` or one list part carrying a string ``text``;
+    images and other parts ride along untouched. Returns None unless the counts
+    line up exactly, so a rewrite never lands on the wrong slot.
+    """
+    if message_text_slot_count(message) != len(texts):
+        return None
+    content: Final = message.get("content")
+    if not isinstance(content, (str, list)):
+        return message
+    rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts)
+    rewritten: Final = {**message, "content": rewritten_content}  # mutable-ok: chat rows stay JSON-plain dicts
+    return cast("AllMessageValues", rewritten)  # cast-ok: the same row with only its text slots swapped
+
+
+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 unappliable_request_rewrite(guardrail_name: str | None) -> UnappliableRequestRewrite:
+    return UnappliableRequestRewrite(guardrail_name or "unknown")

litellm/llms/openai/chat/guardrail_translation/handler.py

@@ -42,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     stream_item_field,
     stream_item_fingerprint,
     stream_item_items,
+    unappliable_request_rewrite,
 )
 from litellm.main import stream_chunk_builder
 from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
@@ -196,6 +197,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
             else:
                 # Step 3: Map guardrail responses back to original message structure
                 if guardrailed_texts and texts_to_check:
+                    if len(guardrailed_texts) != len(text_task_mappings):
+                        raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
                     await self._apply_guardrail_responses_to_input_texts(
                         messages=messages,
                         responses=guardrailed_texts,

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -56,6 +56,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
     stream_item_field,
     stream_item_fingerprint,
     stream_item_items,
+    unappliable_request_rewrite,
 )
 from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
 from litellm.responses.litellm_completion_transformation.transformation import (
@@ -495,13 +496,13 @@ class OpenAIResponsesHandler(BaseTranslation):
                 data["instructions"] = written_back.instructions  # rebind-ok: data is an out-param
         elif isinstance(input_data, str):
             guardrailed_texts: Final = guardrailed_inputs.get("texts") or ()
+            if len(guardrailed_texts) > 1:
+                raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
             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")
+                raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
             await self._apply_guardrail_responses_to_input(
                 messages=input_data,
                 responses=rewritten_texts,

litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py

@@ -7,7 +7,7 @@
 
 import fnmatch
 import os
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional
 
 import httpx
@@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import (
     httpxSpecialProvider,
 )
 from litellm.types.guardrails import GuardrailEventHooks
-from litellm.types.llms.openai import ChatCompletionToolParam
+from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
 from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
     GenericGuardrailAPIMetadata,
     GenericGuardrailAPIRequest,
@@ -150,6 +150,26 @@ def _extract_inbound_headers(
     return None
 
 
+def _structured_rows_to_write_back(
+    original_rows: Sequence[AllMessageValues] | None,
+    shown_rows: Sequence[AllMessageValues] | None,
+    returned_rows: Sequence[AllMessageValues],
+) -> tuple[AllMessageValues, ...] | None:
+    """The request model drops row keys its message types do not declare, so a
+    row the server echoes back verbatim is restored to the original row object.
+    A server that echoes every row back unchanged has not rewritten anything
+    per row, so its answer is read from texts, as it was before rows could be
+    returned at all."""
+    if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows):
+        return tuple(returned_rows)
+    if all(returned == shown for shown, returned in zip(shown_rows, returned_rows)):
+        return None
+    return tuple(
+        original if returned == shown else returned
+        for original, shown, returned in zip(original_rows, shown_rows, returned_rows)
+    )
+
+
 class GenericGuardrailAPI(CustomGuardrail):
     """
     Generic Guardrail API integration for LiteLLM.
@@ -322,6 +342,8 @@ class GenericGuardrailAPI(CustomGuardrail):
         texts: list,
         images: list[str] | None,
         tools: list[ChatCompletionToolParam] | None,
+        structured_messages: Sequence[AllMessageValues] | None,
+        shown_messages: Sequence[AllMessageValues] | None,
         guardrail_response: GenericGuardrailAPIResponse,
     ) -> GenericGuardrailAPIInputs:
         # Action is NONE or no modifications needed
@@ -336,6 +358,13 @@ class GenericGuardrailAPI(CustomGuardrail):
             return_inputs["tools"] = guardrail_response.tools
         elif tools:
             return_inputs["tools"] = tools
+        rows_to_write_back: Final = (
+            _structured_rows_to_write_back(structured_messages, shown_messages, guardrail_response.structured_messages)
+            if guardrail_response.structured_messages
+            else None
+        )
+        if rows_to_write_back is not None:
+            return_inputs["structured_messages"] = list(rows_to_write_back)  # mutable-ok: guardrail inputs take a list
         if guardrail_response.stream_holdback_chars is not None:
             return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars
         return return_inputs
@@ -473,6 +502,8 @@ class GenericGuardrailAPI(CustomGuardrail):
                 texts=texts,
                 images=images,
                 tools=tools,
+                structured_messages=structured_messages,
+                shown_messages=guardrail_request.structured_messages,
                 guardrail_response=guardrail_response,
             )
 

litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py

@@ -2,6 +2,7 @@ import asyncio
 import base64
 import os
 from collections.abc import Mapping, Sequence
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, Literal, Optional
 
 import httpx
@@ -14,11 +15,13 @@ from litellm.integrations.custom_guardrail import (
     CustomGuardrail,
     log_guardrail_information,
 )
+from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts, message_with_slot_texts
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
     httpxSpecialProvider,
 )
 from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import GenericGuardrailAPIInputs
 
 if TYPE_CHECKING:
@@ -28,12 +31,36 @@ if TYPE_CHECKING:
 
 _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
 _SANITIZE_FILE_QUEUED_STATUSES: Final = frozenset({"created", "in progress"})
+_PROTECT_ROLES: Final = frozenset({"system", "user", "assistant"})
 
 
 class PromptSecurityGuardrailMissingSecrets(Exception):
     pass
 
 
+def _inputs_with_structured_messages(
+    inputs: GenericGuardrailAPIInputs, rewritten_messages: Sequence[AllMessageValues] | None
+) -> GenericGuardrailAPIInputs:
+    if rewritten_messages is None:
+        return inputs
+    patched: Final[GenericGuardrailAPIInputs] = {
+        **inputs,
+        "structured_messages": list(rewritten_messages),  # mutable-ok: the TypedDict field is declared as a list
+    }
+    return patched
+
+
+def _inputs_with_modifications(
+    inputs: GenericGuardrailAPIInputs,
+    modified_texts: list[str],
+    rewritten_messages: Sequence[AllMessageValues] | None,
+) -> GenericGuardrailAPIInputs:
+    if not modified_texts:
+        return _inputs_with_structured_messages(inputs, rewritten_messages)
+    with_texts: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": modified_texts}
+    return _inputs_with_structured_messages(with_texts, rewritten_messages)
+
+
 class _ProtectVerdict(TypedDict, total=False):
     """One side (``prompt`` or ``response``) of an ``/api/protect`` verdict."""
 
@@ -276,14 +303,39 @@ class PromptSecurityGuardrail(CustomGuardrail):
                 detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
             )
         elif action == "modify":
-            # Extract modified texts from modified_messages
             modified_messages: Final = result.get("modified_messages", [])
-            modified_texts: Final = self._extract_texts_from_messages(modified_messages)
-            if modified_texts:
-                inputs["texts"] = modified_texts
+            return _inputs_with_modifications(
+                inputs,
+                self._extract_texts_from_messages(modified_messages),
+                self._structured_messages_with_modifications(structured_messages, modified_messages),
+            )
 
         return inputs
 
+    def _is_sent_to_protect(self, message: Mapping[str, object]) -> bool:
+        return self.check_tool_results or message.get("role") in _PROTECT_ROLES
+
+    def _structured_messages_with_modifications(
+        self,
+        structured_messages: Sequence[AllMessageValues],
+        modified_messages: Sequence[Mapping[str, object]],
+    ) -> tuple[AllMessageValues, ...] | None:
+        sent_indices: Final = tuple(
+            index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message)
+        )
+        if not sent_indices or len(sent_indices) != len(modified_messages):
+            return None
+        rewritten: Final = tuple(
+            message_with_slot_texts(structured_messages[index], self._extract_texts_from_messages((modified,)))
+            for index, modified in zip(sent_indices, modified_messages)
+        )
+        replacements: Final = MappingProxyType(
+            {index: message for index, message in zip(sent_indices, rewritten) if message is not None}
+        )
+        if len(replacements) != len(sent_indices):
+            return None
+        return tuple(replacements.get(index, message) for index, message in enumerate(structured_messages))
+
     async def _apply_guardrail_on_response(
         self,
         inputs: GenericGuardrailAPIInputs,
@@ -347,19 +399,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
         return inputs
 
     def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]:
-        """Extract text content from messages."""
-        texts: Final = []
-        for message in messages:
-            content = message.get("content")
-            if isinstance(content, str):
-                texts.append(content)
-            elif isinstance(content, list):
-                for item in content:
-                    if isinstance(item, dict) and item.get("type") == "text":
-                        text = item.get("text")
-                        if text:
-                            texts.append(text)
-        return texts
+        return [text for message in messages for text in message_slot_texts(message)]
 
     async def _process_standalone_images(self, images: list[str], user_api_key_alias: str | None) -> None:
         """Process standalone images from inputs (data URLs)."""
@@ -681,14 +721,13 @@ class PromptSecurityGuardrail(CustomGuardrail):
 
         This allows checking tool results for indirect prompt injection when enabled.
         """
-        supported_roles: Final = ["system", "user", "assistant"]
         filtered_messages: Final = []
         transformed_count = 0
         filtered_count = 0
 
         for message in messages:
             role = message.get("role", "")
-            if role in supported_roles:
+            if role in _PROTECT_ROLES:
                 filtered_messages.append(message)
             else:
                 if self.check_tool_results:

litellm/proxy/policy_engine/pipeline_executor.py

@@ -58,15 +58,6 @@ 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

litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py

@@ -1,4 +1,5 @@
-from typing import Any, Final, Literal
+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 BaseModel, ConfigDict, Field
 from typing_extensions import TypedDict
@@ -158,12 +159,21 @@ def coerce_stream_holdback_value(value: Any) -> int:
         return 0
 
 
+def structured_messages_from_response(value: object) -> Sequence[AllMessageValues] | None:
+    if not isinstance(value, list):
+        return None
+    if not all(isinstance(message, Mapping) and isinstance(message.get("role"), str) for message in value):
+        return None
+    return cast("Sequence[AllMessageValues]", value)  # cast-ok: JSON rows checked for a role, the same trust texts get
+
+
 class GenericGuardrailAPIResponse:
     """Response model for the Generic Guardrail API"""
 
     texts: list[str] | None
     images: list[str] | None
     tools: list[GuardrailToolParam] | None
+    structured_messages: Sequence[AllMessageValues] | None
     action: str
     blocked_reason: str | None
     stream_holdback_chars: list[int] | None
@@ -176,12 +186,14 @@ class GenericGuardrailAPIResponse:
         images: list[str] | None = None,
         tools: list[GuardrailToolParam] | None = None,
         stream_holdback_chars: list[int] | None = None,
+        structured_messages: Sequence[AllMessageValues] | None = None,
     ) -> None:
         self.action = action
         self.blocked_reason = blocked_reason
         self.texts = texts
         self.images = images
         self.tools = tools
+        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
         # processing round (word-boundary safety for text transformations).
@@ -200,4 +212,5 @@ class GenericGuardrailAPIResponse:
             images=data.get("images"),
             tools=data.get("tools"),
             stream_holdback_chars=stream_holdback_chars,
+            structured_messages=structured_messages_from_response(data.get("structured_messages")),
         )

tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py

@@ -2272,6 +2272,58 @@ class TestAnthropicMessagesHandlerStreamingScanKey:
         assert ended_key != open_key
 
 
+class PerRowTextGuardrail(CustomGuardrail):
+    """Answers one redacted text per chat row it was shown, the way a guardrail
+    that scans per message does, and hands back only texts."""
+
+    def __init__(self):
+        super().__init__(guardrail_name="per-row-redactor")
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        rows = inputs.get("structured_messages") or []
+        return {**inputs, "texts": [str(row.get("content")).replace("123-45-6789", "<US_SSN>") for row in rows]}
+
+
+class TestPerMessageTextWriteBack:
+    """Texts that no longer pair one-to-one with what the handler extracted must be
+    rejected by name instead of sliding onto the wrong messages."""
+
+    @pytest.mark.asyncio
+    async def test_one_text_per_row_over_a_system_prompt_is_rejected_by_name(self):
+        from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
+
+        data = {
+            "model": "claude-sonnet-4-5",
+            "system": "Reply with exactly the SSN you were given.",
+            "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
+        }
+        original = json.loads(json.dumps(data))
+
+        with pytest.raises(UnappliableRequestRewrite) as excinfo:
+            await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail())
+
+        assert excinfo.value.guardrail_name == "per-row-redactor"
+        assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched"
+        assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched"
+
+    @pytest.mark.asyncio
+    async def test_one_text_per_row_without_a_system_prompt_is_applied(self):
+        data = {
+            "model": "claude-sonnet-4-5",
+            "messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
+        }
+
+        await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail())
+
+        assert data["messages"] == [{"role": "user", "content": "My SSN is <US_SSN>."}]
+
+
 class TestAnthropicMessagesHandlerPostCallHookResponse:
     def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self):
         from litellm.types.utils import Choices, Message, ModelResponse, Usage

tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py

@@ -1893,6 +1893,49 @@ class TestScanOnlyToolResults:
         assert data["messages"][4]["content"] == "and then?"
 
 
+class ToolDroppingTextGuardrail(CustomGuardrail):
+    """Answers one text per non-tool message it saw, the way a guardrail that
+    filters tool rows out before scanning does, and hands back only texts."""
+
+    def __init__(self):
+        super().__init__(guardrail_name="tool-dropping-redactor")
+
+    async def apply_guardrail(
+        self,
+        inputs: GenericGuardrailAPIInputs,
+        request_data: dict,
+        input_type: Literal["request", "response"],
+        logging_obj: Optional[Any] = None,
+    ) -> GenericGuardrailAPIInputs:
+        kept = [m for m in inputs.get("structured_messages") or [] if m.get("role") != "tool"]
+        return {**inputs, "texts": [str(m.get("content")).replace("POISON", "[BLOCKED]") for m in kept]}
+
+
+class TestPerMessageTextWriteBack:
+    """Texts that no longer pair one-to-one with what the handler extracted must be
+    rejected by name instead of sliding onto the wrong messages."""
+
+    @pytest.mark.asyncio
+    async def test_fewer_texts_than_extracted_over_a_tool_message_is_rejected(self):
+        from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
+
+        handler = OpenAIChatCompletionsHandler()
+        original_messages = [
+            {"role": "system", "content": "SYSTEM-PROMPT"},
+            {"role": "user", "content": "fetch the page"},
+            {"role": "assistant", "content": "fetching"},
+            {"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"},
+            {"role": "user", "content": "and then?"},
+        ]
+        data = {"messages": json.loads(json.dumps(original_messages))}
+
+        with pytest.raises(UnappliableRequestRewrite) as excinfo:
+            await handler.process_input_messages(data=data, guardrail_to_apply=ToolDroppingTextGuardrail())
+
+        assert excinfo.value.guardrail_name == "tool-dropping-redactor"
+        assert data["messages"] == original_messages, "a rejected rewrite must leave the request untouched"
+
+
 class TestBuildBlockSseChunks:
     """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks"""
 

tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py

@@ -8,7 +8,7 @@ with guardrail transformations.
 import copy
 from collections.abc import Callable
 from typing import Any, List, Literal, Optional, Tuple
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import logging
 
@@ -31,6 +31,7 @@ from litellm.llms.openai.responses.guardrail_translation.handler import (
     OpenAIResponsesHandler,
 )
 from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
+from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI
 from litellm.types.llms.openai import ChatCompletionToolCallChunk
 from litellm.responses.litellm_completion_transformation.transformation import (
     LiteLLMCompletionResponsesConfig,
@@ -2338,6 +2339,135 @@ def _parallel_tool_call_input() -> list:
     ]
 
 
+SSN = "123-45-6789"
+REDACTED_SSN = "<US_SSN>"
+
+
+def _redacted(value: object) -> object:
+    if isinstance(value, str):
+        return value.replace(SSN, REDACTED_SSN)
+    if isinstance(value, list):
+        return [{**part, "text": _redacted(part["text"])} if "text" in part else part for part in value]
+    return value
+
+
+def _per_message_guardrail_server(structured_messages_in_answer: bool) -> Callable[..., MagicMock]:
+    """Answers one redacted text per chat row it was shown, the way a guardrail
+    that scans per message does, and optionally the rewritten rows themselves."""
+
+    def post(url: str, json: dict, headers: dict) -> MagicMock:
+        rows = json["structured_messages"]
+        answer: dict = {
+            "action": "GUARDRAIL_INTERVENED",
+            "texts": [_redacted(row["content"]) if isinstance(row.get("content"), str) else "" for row in rows],
+        }
+        if structured_messages_in_answer:
+            answer["structured_messages"] = [{**row, "content": _redacted(row.get("content"))} for row in rows]
+        response = MagicMock()
+        response.json.return_value = answer
+        response.raise_for_status = MagicMock()
+        return response
+
+    return post
+
+
+def _per_message_redactor() -> GenericGuardrailAPI:
+    return GenericGuardrailAPI(
+        api_base="https://guardrail.test",
+        guardrail_name="per-message-redactor",
+        event_hook="pre_call",
+        default_on=True,
+    )
+
+
+def _tool_replay_request() -> dict:
+    return {
+        "model": "gpt-5.6",
+        "instructions": "Never repeat the SSN " + SSN + " back.",
+        "input": [
+            {"role": "user", "content": "Look up " + SSN + " for me."},
+            {"type": "function_call", "call_id": "call_1", "name": "lookup_customer", "arguments": '{"id": "42"}'},
+            {"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'},
+        ],
+    }
+
+
+def _string_input_request() -> dict:
+    return {
+        "model": "gpt-5.6",
+        "instructions": "Never repeat the SSN " + SSN + " back.",
+        "input": "My SSN is " + SSN + ".",
+    }
+
+
+class TestPerMessageRewriteWriteBack:
+    """A guardrail that rewrites per chat row hands the rows back as
+    structured_messages, and the handler lands them on the instructions and the
+    input items they came from; the same rewrite handed back as texts alone has
+    no item to land on and is rejected by name instead of sent unrewritten."""
+
+    @pytest.mark.asyncio
+    async def test_structured_rows_land_on_instructions_and_tool_output(self):
+        guardrail = _per_message_redactor()
+        data = _tool_replay_request()
+        function_call_item = data["input"][1]
+
+        with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)):
+            result = await OpenAIResponsesHandler().process_input_messages(data, guardrail)
+
+        assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back."
+        assert _texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."]
+        assert result["input"][1] == function_call_item
+        assert result["input"][2] == {
+            "type": "function_call_output",
+            "call_id": "call_1",
+            "output": '{"ssn": "' + REDACTED_SSN + '"}',
+        }
+
+    @pytest.mark.asyncio
+    async def test_texts_only_per_message_answer_is_rejected_by_name(self):
+        from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
+
+        guardrail = _per_message_redactor()
+        data = _tool_replay_request()
+        original = copy.deepcopy(data)
+
+        with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)):
+            with pytest.raises(UnappliableRequestRewrite) as excinfo:
+                await OpenAIResponsesHandler().process_input_messages(data, guardrail)
+
+        assert excinfo.value.guardrail_name == "per-message-redactor"
+        assert data["input"] == original["input"]
+        assert data["instructions"] == original["instructions"]
+
+    @pytest.mark.asyncio
+    async def test_structured_rows_land_on_instructions_and_string_input(self):
+        guardrail = _per_message_redactor()
+        data = _string_input_request()
+
+        with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)):
+            result = await OpenAIResponsesHandler().process_input_messages(data, guardrail)
+
+        assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back."
+        assert [_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]]
+
+    @pytest.mark.asyncio
+    async def test_texts_only_per_message_answer_over_a_string_input_is_rejected_by_name(self):
+        from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
+
+        guardrail = _per_message_redactor()
+        data = _string_input_request()
+        original = copy.deepcopy(data)
+
+        with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)):
+            with pytest.raises(UnappliableRequestRewrite) as excinfo:
+                await OpenAIResponsesHandler().process_input_messages(data, guardrail)
+
+        assert excinfo.value.guardrail_name == "per-message-redactor"
+        assert data["input"] == original["input"]
+        assert data["instructions"] == original["instructions"]
+
+
 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."""

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py

@@ -1820,7 +1820,7 @@ async def test_unalignable_rewrite_is_rejected_never_sent_unredacted(
     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
+    from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
 
     data: dict[str, object] = {"model": "gpt-4o", "input": responses_input}
     if instructions is not None:

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py

@@ -582,6 +582,145 @@ class TestGuardrailActions:
             assert result_images is None
 
 
+class TestStructuredMessagesInResponse:
+    """A guardrail server that rewrites per chat row answers with the rewritten
+    rows as structured_messages, which the endpoint handlers write back by row."""
+
+    @pytest.mark.asyncio
+    async def test_returned_rows_are_handed_back_as_structured_messages(
+        self, generic_guardrail, mock_request_data_input
+    ):
+        rewritten_rows = [
+            {"role": "system", "content": "Never repeat an SSN."},
+            {"role": "user", "content": "Look up [REDACTED] for me."},
+            {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'},
+        ]
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "action": "GUARDRAIL_INTERVENED",
+            "texts": ["Never repeat an SSN.", "Look up [REDACTED] for me.", '{"ssn": "[REDACTED]"}'],
+            "structured_messages": rewritten_rows,
+        }
+        mock_response.raise_for_status = MagicMock()
+
+        with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response):
+            guardrailed_inputs = await generic_guardrail.apply_guardrail(
+                inputs={"texts": ["Look up 123-45-6789 for me."]},
+                request_data=mock_request_data_input,
+                input_type="request",
+            )
+
+        assert guardrailed_inputs["structured_messages"] == rewritten_rows
+        assert guardrailed_inputs["texts"] == mock_response.json.return_value["texts"]
+
+    @pytest.mark.asyncio
+    async def test_rows_echoed_back_as_shown_keep_their_original_keys(
+        self, generic_guardrail, mock_request_data_input
+    ):
+        tool_call_row = {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}, "index": 0}
+            ],
+        }
+        original_rows = [
+            {"role": "user", "content": "Look up 123-45-6789 for me.", "name": "pat"},
+            tool_call_row,
+            {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "123-45-6789"}'},
+        ]
+
+        def echo_with_tool_output_redacted(url, json, headers):
+            shown_rows = json["structured_messages"]
+            assert "index" not in shown_rows[1]["tool_calls"][0]
+            assert "name" not in shown_rows[0]
+            answer = MagicMock()
+            answer.json.return_value = {
+                "action": "GUARDRAIL_INTERVENED",
+                "texts": ["Look up 123-45-6789 for me."],
+                "structured_messages": [
+                    shown_rows[0],
+                    shown_rows[1],
+                    {**shown_rows[2], "content": '{"ssn": "[REDACTED]"}'},
+                ],
+            }
+            answer.raise_for_status = MagicMock()
+            return answer
+
+        with patch.object(generic_guardrail.async_handler, "post", side_effect=echo_with_tool_output_redacted):
+            guardrailed_inputs = await generic_guardrail.apply_guardrail(
+                inputs={"texts": ["Look up 123-45-6789 for me."], "structured_messages": original_rows},
+                request_data=mock_request_data_input,
+                input_type="request",
+            )
+
+        returned_rows = guardrailed_inputs["structured_messages"]
+        assert returned_rows[0] is original_rows[0]
+        assert returned_rows[1] is tool_call_row
+        assert returned_rows[2] == {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "[REDACTED]"}'}
+
+    @pytest.mark.asyncio
+    async def test_rows_all_echoed_back_as_shown_leave_the_rewrite_to_texts(
+        self, generic_guardrail, mock_request_data_input
+    ):
+        """A server written against the texts contract that echoes the request rows back
+        untouched while rewriting texts still gets its texts rewrite applied."""
+        original_rows = [
+            {"role": "system", "content": "Never repeat an SSN."},
+            {"role": "user", "content": "Look up 123-45-6789 for me."},
+        ]
+
+        def echo_rows_and_rewrite_texts(url, json, headers):
+            answer = MagicMock()
+            answer.json.return_value = {
+                "action": "NONE",
+                "texts": [text.replace("123-45-6789", "[REDACTED]") for text in json["texts"]],
+                "structured_messages": json["structured_messages"],
+            }
+            answer.raise_for_status = MagicMock()
+            return answer
+
+        with patch.object(generic_guardrail.async_handler, "post", side_effect=echo_rows_and_rewrite_texts):
+            guardrailed_inputs = await generic_guardrail.apply_guardrail(
+                inputs={
+                    "texts": ["Never repeat an SSN.", "Look up 123-45-6789 for me."],
+                    "structured_messages": original_rows,
+                },
+                request_data=mock_request_data_input,
+                input_type="request",
+            )
+
+        assert "structured_messages" not in guardrailed_inputs
+        assert guardrailed_inputs["texts"] == ["Never repeat an SSN.", "Look up [REDACTED] for me."]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "structured_messages",
+        [[], [{"content": "a row with no role"}], "not a list"],
+        ids=["empty", "no_role", "not_a_list"],
+    )
+    async def test_rows_that_are_not_chat_messages_are_ignored(
+        self, generic_guardrail, mock_request_data_input, structured_messages
+    ):
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "action": "GUARDRAIL_INTERVENED",
+            "texts": ["[REDACTED]"],
+            "structured_messages": structured_messages,
+        }
+        mock_response.raise_for_status = MagicMock()
+
+        with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response):
+            guardrailed_inputs = await generic_guardrail.apply_guardrail(
+                inputs={"texts": ["Look up 123-45-6789 for me."]},
+                request_data=mock_request_data_input,
+                input_type="request",
+            )
+
+        assert "structured_messages" not in guardrailed_inputs
+        assert guardrailed_inputs["texts"] == ["[REDACTED]"]
+
+
 class TestImageSupport:
     """Test image handling in guardrail requests"""
 

tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py

@@ -1,5 +1,6 @@
 import asyncio
 import base64
+from collections.abc import Mapping, Sequence
 from unittest.mock import AsyncMock, patch
 
 import pytest
@@ -12,6 +13,7 @@ from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security im
     PromptSecurityGuardrailMissingSecrets,
 )
 from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
+from litellm.types.llms.openai import AllMessageValues
 
 
 def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch):
@@ -174,6 +176,123 @@ async def test_apply_guardrail_modify_request(monkeypatch: pytest.MonkeyPatch):
     assert result["texts"] == ["User prompt with PII: SSN [REDACTED]"]
 
 
+def _modify_response(modified_messages: Sequence[Mapping[str, object]]) -> Response:
+    mock_response = Response(
+        json={"result": {"prompt": {"action": "modify", "modified_messages": modified_messages}}},
+        status_code=200,
+        request=Request(method="POST", url="https://test.prompt.security/api/protect"),
+    )
+    mock_response.raise_for_status = lambda: None
+    return mock_response
+
+
+def _tool_replay_messages() -> list[AllMessageValues]:
+    return [
+        {"role": "system", "content": "Never echo an SSN like 123-45-6789."},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Look up 123-45-6789"},
+                {"type": "image_url", "image_url": {"url": "https://example.com/id-card.png"}},
+            ],
+        },
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}
+            ],
+        },
+        {"role": "tool", "tool_call_id": "call_1", "content": '{"ssn": "123-45-6789"}'},
+        {"role": "user", "content": "Summarize what you found."},
+    ]
+
+
+@pytest.mark.asyncio
+async def test_modify_returns_structured_messages_with_tool_rows_kept(monkeypatch: pytest.MonkeyPatch):
+    """A per-message modify verdict comes back as structured_messages so the
+    endpoint handler can write it back by message, with the rows Prompt Security
+    never saw (tool results) and the non-text parts (images) left in place."""
+    monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
+    monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
+    guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True)
+    messages = _tool_replay_messages()
+    inputs = {"texts": ["Look up 123-45-6789", "Summarize what you found."], "structured_messages": messages}
+    modified_messages = [
+        {"role": "system", "content": "Never echo an SSN like [REDACTED]."},
+        {"role": "user", "content": [{"type": "text", "text": "Look up [REDACTED]"}]},
+        {"role": "assistant", "content": None},
+        {"role": "user", "content": "Summarize what you found."},
+    ]
+
+    with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)):
+        result = await guardrail.apply_guardrail(
+            inputs=inputs, request_data={"messages": messages}, input_type="request"
+        )
+
+    assert result["structured_messages"] == [
+        {"role": "system", "content": "Never echo an SSN like [REDACTED]."},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Look up [REDACTED]"},
+                {"type": "image_url", "image_url": {"url": "https://example.com/id-card.png"}},
+            ],
+        },
+        messages[2],
+        messages[3],
+        {"role": "user", "content": "Summarize what you found."},
+    ]
+    assert result["structured_messages"] is not messages
+    assert result["texts"] == [
+        "Never echo an SSN like [REDACTED].",
+        "Look up [REDACTED]",
+        "Summarize what you found.",
+    ]
+
+
+@pytest.mark.asyncio
+async def test_modify_with_unexpected_message_count_keeps_texts_only(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
+    monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
+    guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True)
+    messages = _tool_replay_messages()
+    inputs = {"texts": ["Look up 123-45-6789", "Summarize what you found."], "structured_messages": messages}
+    modified_messages = [{"role": "user", "content": "Look up [REDACTED]"}]
+
+    with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)):
+        result = await guardrail.apply_guardrail(
+            inputs=inputs, request_data={"messages": messages}, input_type="request"
+        )
+
+    assert result["structured_messages"] is messages
+    assert result["texts"] == ["Look up [REDACTED]"]
+
+
+@pytest.mark.asyncio
+async def test_modify_keeps_empty_text_parts_as_slots(monkeypatch: pytest.MonkeyPatch):
+    """The chat handler counts an empty text part as a slot, so a modify verdict
+    that echoes the empty part still lines up with the row and its texts."""
+    monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key")
+    monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security")
+    guardrail = PromptSecurityGuardrail(guardrail_name="test-guard", event_hook="pre_call", default_on=True)
+    messages: list[AllMessageValues] = [
+        {"role": "user", "content": [{"type": "text", "text": "Look up 123-45-6789"}, {"type": "text", "text": ""}]}
+    ]
+    inputs = {"texts": ["Look up 123-45-6789", ""], "structured_messages": messages}
+    modified_messages = [
+        {"role": "user", "content": [{"type": "text", "text": "Look up [REDACTED]"}, {"type": "text", "text": ""}]}
+    ]
+
+    with patch.object(guardrail.async_handler, "post", return_value=_modify_response(modified_messages)):
+        result = await guardrail.apply_guardrail(
+            inputs=inputs, request_data={"messages": messages}, input_type="request"
+        )
+
+    assert result["structured_messages"] == modified_messages
+    assert result["texts"] == ["Look up [REDACTED]", ""]
+
+
 @pytest.mark.asyncio
 async def test_apply_guardrail_allow_request(monkeypatch: pytest.MonkeyPatch):
     """Test that apply_guardrail allows safe prompts"""