fix(guardrails): Propagate mutated text to proxied request/response

plugins/nemo-guardrails/src/nemo_guardrails_plugin/middleware.py

@@ -44,6 +44,7 @@
 )
 from nemo_guardrails_plugin.responses import (
     GUARDRAILS_DATA_FIELD,
+    apply_input_rail_modifications,
     build_assistant_message_from_response_result,
     build_blocked_immediate_response_body,
     build_blocked_output_response_body,
@@ -339,10 +340,10 @@ async def process_request(
 
         # Remove Guardrails-specific fields from the request body proxied to the upstream model.
         # Otherwise, the upstream model may reject the request.
-        # Nested values inside ``request.body`` (notably ``messages``) remain aliased with
-        # ``ctx.original_request.body`` per IGW's shallow snapshot — fine because we don't mutate them.
-        sanitized_body = sanitize_request_body_for_proxy(request.body)
-        request = InferenceRequest(body=sanitized_body, headers=request.headers, path=request.path)
+        # Assign a new top-level body dict (sanitize returns a shallow copy) so we do not
+        # mutate ``ctx.original_request.body``. Nested ``messages`` may still be aliased
+        # until replaced below — we only replace that list, never mutate it in place.
+        request.body = sanitize_request_body_for_proxy(request.body)
 
         if not source_has_input_flows(source):
             return request
@@ -373,6 +374,13 @@ async def process_request(
                 )
             )
 
+        # If an input rail masked/transformed the user message, write the text content
+        # back onto the last user message so the upstream model sees the post-rail content.
+        updated_messages = apply_input_rail_modifications(messages, generation_response)
+        if updated_messages is not messages:
+            request.body["messages"] = updated_messages
+            logger.debug("Applied input-rail message modifications for %s", provenance.label)
+
         # Store the generation_response in plugin state so it can be used by the response middleware
         # to build the `guardrails_data` for the input and output rails.
         logger.debug("Storing process_request GenerationResponse for %s", provenance.label)

plugins/nemo-guardrails/src/nemo_guardrails_plugin/responses.py

@@ -144,8 +144,12 @@ def is_blocked_generation_response(generation_response: GenerationResponse) -> b
 
 
 def extract_response_content(generation_response: GenerationResponse) -> str:
-    """
-    Extract the last assistant message content from a GenerationResponse.
+    """Return the post-rail text from a ``GenerationResponse``.
+
+    ``GenerationResponse.response`` is OpenAI-shaped. When it is a message list,
+    the library puts the rails result on an ``role=assistant`` entry — even for
+    **input-only** runs, where that string is the post-rail ``$user_message``
+    (e.g. PII-redacted user text), not a model reply from chat history.
     """
     response = generation_response.response
     if isinstance(response, list):
@@ -158,6 +162,77 @@ def extract_response_content(generation_response: GenerationResponse) -> str:
     return response
 
 
+# ---------------------------------------------------------------------------
+# Masking write-back helpers
+#
+# After input-only generate_async, the (possibly redacted) user text is returned
+# as GenerationResponse.response with role=assistant (not role=user).
+# Example: request user "Hi John" → response [{"role": "assistant", "content": "Hi <PERSON>"}].
+# If the message content is redacted, we need to write it back onto the last user message
+# in the request.
+# ---------------------------------------------------------------------------
+
+
+def _index_of_last_user_message(messages: list[dict[str, Any]]) -> int | None:
+    """Return the index of the last ``role=user`` message, if any."""
+    last_user_index: int | None = None
+    for index, message in enumerate(messages):
+        if isinstance(message, dict) and message.get("role") == "user":
+            last_user_index = index
+    return last_user_index
+
+
+def apply_input_rail_modifications(
+    messages: list[dict[str, Any]],
+    generation_response: GenerationResponse,
+) -> list[dict[str, Any]]:
+    """Return messages with input-rail modifications applied to the last user turn.
+
+    Post-rail user text is taken from ``GenerationResponse.response`` via
+    ``extract_response_content``. If modified, it is written onto the last user message.
+
+    Returns the same list object when there is nothing to change; otherwise, returns a
+    shallow copy that replaces only the last user message dict.
+    """
+    last_user_index = _index_of_last_user_message(messages)
+    if last_user_index is None:
+        return messages
+
+    last_user_message = messages[last_user_index]
+    original_content = last_user_message.get("content")
+    if not isinstance(original_content, str):
+        logger.debug(
+            "Skipping input-rail write-back; last user content is %s, not str",
+            type(original_content).__name__,
+        )
+        return messages
+
+    # NOTE: the input-rail result is labeled as an assistant message in the GenerationResponse,
+    # but the content is actually the redacted/transformed user message.
+    processed = extract_response_content(generation_response)
+    # If the message content wasn't modified, return the original messages.
+    if not isinstance(processed, str) or processed == original_content:
+        return messages
+
+    # If the message content was modified, write it back onto the last user message.
+    updated_messages = list(messages)
+    updated_messages[last_user_index] = {**last_user_message, "content": processed}
+    return updated_messages
+
+
+def apply_output_rail_modifications(
+    original_content: str,
+    generation_response: GenerationResponse,
+) -> str:
+    """Return assistant content after output rails, from ``response`` content."""
+    processed = extract_response_content(generation_response)
+
+    if not isinstance(processed, str) or processed == original_content:
+        return original_content
+
+    return processed
+
+
 def build_assistant_message_from_response_result(response_result: ResponseResult) -> dict[str, Any]:
     """
     Build an assistant message object with the content from the given response.
@@ -299,9 +374,19 @@ def build_output_response_body(
 
     choices = list(response.get("choices", [])) if isinstance(response.get("choices"), list) else []
     response["choices"] = choices
-    # Output rails validate choices[0].message, so return only the choice that was checked
+
+    # Output rails validate choices[0].message, so return only the choice that was checked.
     if generation_response is not None and choices:
-        response["choices"] = [{**choices[0], "index": 0}]
+        choice = {**choices[0], "index": 0}
+        message = dict(choice.get("message") or {}) if isinstance(choice.get("message"), dict) else {}
+        original_content = message.get("content")
+
+        # If the output rail modified the message content, write it back onto the choice.
+        if isinstance(original_content, str):
+            message["content"] = apply_output_rail_modifications(original_content, generation_response)
+            choice["message"] = message
+
+        response["choices"] = [choice]
 
     guardrails_data = build_guardrails_data(
         config_id,

plugins/nemo-guardrails/tests/unit/test_middleware.py

@@ -162,9 +162,13 @@ def _make_entity(
     )
 
 
-def _make_generation_response(*, is_blocked: bool = False) -> GenerationResponse:
+def _make_generation_response(
+    *,
+    is_blocked: bool = False,
+    content: str = "I'm sorry, I can't help with that.",
+) -> GenerationResponse:
     return GenerationResponse(
-        response=[{"role": "assistant", "content": "I'm sorry, I can't help with that."}],
+        response=[{"role": "assistant", "content": content}],
         log=GenerationLog(
             activated_rails=[ActivatedRail(type="input", name="self check input", stop=is_blocked)],
             stats=GenerationStats(input_rails_duration=0.1, total_duration=0.1),
@@ -615,7 +619,7 @@ async def test_invalid_messages_short_circuits(
 
     async def test_successful_generation_returns_request_body(self, middleware: GuardrailsMiddleware) -> None:
         request_body = {"model": "ws/llama", "messages": [{"role": "user", "content": "Hello"}]}
-        generation_response = _make_generation_response(is_blocked=False)
+        generation_response = _make_generation_response(is_blocked=False, content="Hello")
         ctx = _make_ctx(request_body)
 
         with patch.object(middleware, "_run_rails", new=AsyncMock(return_value=generation_response)):
@@ -629,6 +633,54 @@ async def test_successful_generation_returns_request_body(self, middleware: Guar
         assert ctx.state(PLUGIN_NAME).get(STATE_KEY_INPUT_GENERATION_RESPONSE) is generation_response
         assert ctx.response_body_annotations["guardrails_data"]["config_ids"] == ["my-workspace/my-config"]
 
+    async def test_input_masking_writes_back_last_user_message(self, middleware: GuardrailsMiddleware) -> None:
+        request_body = {
+            "model": "ws/llama",
+            "messages": [
+                {"role": "user", "content": "earlier turn"},
+                {"role": "assistant", "content": "ok"},
+                {"role": "user", "content": "Hi! I am Mr. John!"},
+            ],
+        }
+        generation_response = _make_generation_response(
+            is_blocked=False,
+            content="Hi! I am <TITLE> <PERSON>!",
+        )
+
+        with patch.object(middleware, "_run_rails", new=AsyncMock(return_value=generation_response)):
+            result = await _process_request(middleware, request_body, {}, _entity_source())
+
+        assert isinstance(result, dict)
+        assert result["messages"] == [
+            {"role": "user", "content": "earlier turn"},
+            {"role": "assistant", "content": "ok"},
+            {"role": "user", "content": "Hi! I am <TITLE> <PERSON>!"},
+        ]
+        # Original request body must not be mutated (IGW aliases nested values).
+        assert request_body["messages"][-1]["content"] == "Hi! I am Mr. John!"
+
+    async def test_input_masking_skips_non_string_user_content(self, middleware: GuardrailsMiddleware) -> None:
+        # Multimodal content is unsupported for PII write-back; leave the request alone
+        # even when response.content looks like a redacted string (or a stringified list).
+        multimodal_content = [
+            {"type": "text", "text": "Hi! I am Mr. John!"},
+            {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}},
+        ]
+        request_body = {
+            "model": "ws/llama",
+            "messages": [{"role": "user", "content": multimodal_content}],
+        }
+        generation_response = _make_generation_response(
+            is_blocked=False,
+            content=str(multimodal_content),
+        )
+
+        with patch.object(middleware, "_run_rails", new=AsyncMock(return_value=generation_response)):
+            result = await _process_request(middleware, request_body, {}, _entity_source())
+
+        assert isinstance(result, dict)
+        assert result["messages"][0]["content"] == multimodal_content
+
     async def test_user_log_options_forwarded_to_run_rails(self, middleware: GuardrailsMiddleware) -> None:
         request_body = {
             "model": "ws/llama",

plugins/nemo-guardrails/tests/unit/test_responses.py

@@ -9,6 +9,8 @@
 import pytest
 from nemo_guardrails_plugin.constants import GUARDRAILS_DATA_MESSAGE_ROLE
 from nemo_guardrails_plugin.responses import (
+    apply_input_rail_modifications,
+    apply_output_rail_modifications,
     build_assistant_message_from_response_result,
     build_blocked_output_response_body,
     build_immediate_response,
@@ -25,12 +27,18 @@
 # ---------------------------------------------------------------------------
 
 
-def _make_generation_response(*, stopped: bool = False, content: str = "I can't help with that.") -> GenerationResponse:
+def _make_generation_response(
+    *,
+    stopped: bool = False,
+    content: str = "I can't help with that.",
+    output_data: dict[str, Any] | None = None,
+) -> GenerationResponse:
     return GenerationResponse(
         response=[{"role": "assistant", "content": content}],
         log=GenerationLog(
             activated_rails=[ActivatedRail(type="output", name="self check output", stop=stopped)],
         ),
+        output_data=output_data,
     )
 
 
@@ -50,6 +58,75 @@ def _make_response_result(content: str = "Hello!") -> dict[str, Any]:
     }
 
 
+# ---------------------------------------------------------------------------
+# Masking write-back helpers
+# ---------------------------------------------------------------------------
+
+
+class TestApplyInputRailModifications:
+    def test_writes_back_last_user_message_only(self) -> None:
+        earlier_user = {"role": "user", "content": "Hi, I am Alice"}
+        assistant = {"role": "assistant", "content": "Hello Alice"}
+        last_user = {"role": "user", "content": "Hi John"}
+        messages = [earlier_user, assistant, last_user]
+        # Post-rail text comes from GenerationResponse.response content.
+        generation_response = _make_generation_response(content="Hi <PERSON>")
+
+        updated = apply_input_rail_modifications(messages, generation_response)
+
+        assert updated is not messages
+        assert updated[0] is earlier_user
+        assert updated[1] is assistant
+        assert updated[0]["content"] == "Hi, I am Alice"
+        assert updated[1]["content"] == "Hello Alice"
+        assert updated[2]["content"] == "Hi <PERSON>"
+        assert updated[2] is not last_user
+        # Original request messages must not be mutated.
+        assert last_user["content"] == "Hi John"
+
+    def test_returns_same_list_for_non_string_last_user_content(self) -> None:
+        # Multimodal: never write back, even if response.content is a string.
+        messages = [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Hi John"}],
+            }
+        ]
+        generation_response = _make_generation_response(content="Hi <PERSON>")
+
+        assert apply_input_rail_modifications(messages, generation_response) is messages
+
+    def test_returns_same_list_when_content_unchanged(self) -> None:
+        messages = [{"role": "user", "content": "hello"}]
+        generation_response = _make_generation_response(content="hello")
+
+        assert apply_input_rail_modifications(messages, generation_response) is messages
+
+    def test_skips_stringified_multimodal_response_content(self) -> None:
+        # No-op multimodal rails may stringify the list into response.content.
+        # Original content is still a list, so write-back must not run.
+        multimodal_content = [
+            {"type": "text", "text": "hello"},
+            {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAA"}},
+        ]
+        messages = [{"role": "user", "content": multimodal_content}]
+        generation_response = _make_generation_response(content=str(multimodal_content))
+
+        assert apply_input_rail_modifications(messages, generation_response) is messages
+
+
+class TestApplyOutputRailModifications:
+    def test_writes_back_from_response_content(self) -> None:
+        generation_response = _make_generation_response(content="Hi <PERSON>")
+
+        assert apply_output_rail_modifications("Hi John", generation_response) == "Hi <PERSON>"
+
+    def test_returns_original_when_unchanged(self) -> None:
+        generation_response = _make_generation_response(content="hello")
+
+        assert apply_output_rail_modifications("hello", generation_response) == "hello"
+
+
 # ---------------------------------------------------------------------------
 # build_assistant_message_from_response_result
 # ---------------------------------------------------------------------------
@@ -169,7 +246,7 @@ def test_preserves_single_choice_sets_guardrails_data(self) -> None:
         result = build_output_response_body(
             config_id="ws/my-config",
             original_response=original,
-            generation_response=_make_generation_response(),
+            generation_response=_make_generation_response(content="Hello!"),
             input_generation_response=None,
             user_log_options=None,
         )
@@ -178,6 +255,22 @@ def test_preserves_single_choice_sets_guardrails_data(self) -> None:
         assert "guardrails_data" in result
         assert result["guardrails_data"]["config_ids"] == ["ws/my-config"]
 
+    def test_applies_output_masking_to_assistant_content(self) -> None:
+        original = _make_response_result("Hello there! My name is Michael!")
+
+        result = build_output_response_body(
+            config_id="ws/my-config",
+            original_response=original,
+            generation_response=_make_generation_response(
+                content="Hello there! My name is <PERSON>!",
+            ),
+            input_generation_response=None,
+            user_log_options=None,
+        )
+
+        assert result["choices"][0]["message"]["content"] == "Hello there! My name is <PERSON>!"
+        assert original["choices"][0]["message"]["content"] == "Hello there! My name is Michael!"
+
     def test_keeps_only_first_choice(self) -> None:
         original = {
             "id": "chatcmpl-123",
@@ -190,7 +283,7 @@ def test_keeps_only_first_choice(self) -> None:
         result = build_output_response_body(
             config_id="ws/my-config",
             original_response=original,
-            generation_response=_make_generation_response(),
+            generation_response=_make_generation_response(content="A"),
             input_generation_response=None,
             user_log_options=None,
         )
@@ -211,7 +304,7 @@ def test_return_choice_appends_at_correct_index(self) -> None:
         result = build_output_response_body(
             config_id="ws/my-config",
             original_response=original,
-            generation_response=_make_generation_response(),
+            generation_response=_make_generation_response(content="A"),
             input_generation_response=None,
             user_log_options=None,
             return_guardrails_data_as_choice=True,