fix(middleware): handle MODIFIED status in GuardrailsMiddleware instead of silently dropping it (#1714)

Author: Pouyanpi

docs/integration/langchain/agent-middleware.md

@@ -326,9 +326,9 @@ guardrails = GuardrailsMiddleware(
 
 Rails evaluate the `content` field of messages, not the `tool_calls` arguments. Content-based rails do not inspect PII or harmful content passed through tool call arguments (e.g., `send_email(body="SSN: 123-45-6789")`).
 
-### MODIFIED Status Is Ignored
+### MODIFIED Status Replaces Message Content
 
-When a rail modifies content (returns `RailStatus.MODIFIED`), the middleware treats it as a pass-through and the agent uses the original, unmodified content. This is by design — applying modifications to the agent's internal state could cause inconsistencies.
+When a rail modifies content (returns `RailStatus.MODIFIED`), the middleware replaces the relevant message with the modified content. For input rails, the last user message is replaced. For output rails, the last AI message is replaced. This enables use cases like PII redaction and content sanitization.
 
 ---
 

nemoguardrails/integrations/langchain/middleware.py

@@ -111,6 +111,10 @@ async def abefore_model(self, state: AgentState, runtime: LangGraphRuntime) -> O
         if not messages:
             return None
 
+        last_user_message = self._get_last_user_message(messages)
+        if not last_user_message:
+            return None
+
         rails_messages = self._convert_to_rails_messages(messages)
 
         try:
@@ -125,6 +129,11 @@ async def abefore_model(self, state: AgentState, runtime: LangGraphRuntime) -> O
                 blocked_msg = create_ai_message(self.blocked_input_message)
                 return {"messages": messages + [blocked_msg], "jump_to": "end"}
 
+            if result.status == RailStatus.MODIFIED:
+                log.info("Input modified by rail '%s': content replaced", result.rail or "unknown rail")
+                modified_msg = last_user_message.model_copy(update={"content": result.content})
+                return {"messages": self._replace_last_human_message(messages, modified_msg)}
+
             return None
 
         except GuardrailViolation:
@@ -141,6 +150,12 @@ async def abefore_model(self, state: AgentState, runtime: LangGraphRuntime) -> O
             blocked_msg = create_ai_message(self.blocked_input_message)
             return {"messages": messages + [blocked_msg], "jump_to": "end"}
 
+    def _replace_last_human_message(self, messages: list, replacement: HumanMessage) -> list:
+        for i in range(len(messages) - 1, -1, -1):
+            if is_human_message(messages[i]):
+                return messages[:i] + [replacement] + messages[i + 1 :]
+        return messages + [replacement]
+
     def _replace_last_ai_message(self, messages: list, replacement: AIMessage) -> list:
         for i in range(len(messages) - 1, -1, -1):
             if is_ai_message(messages[i]):
@@ -173,6 +188,11 @@ async def aafter_model(self, state: AgentState, runtime: LangGraphRuntime) -> Op
                 blocked_msg = create_ai_message(self.blocked_output_message)
                 return {"messages": self._replace_last_ai_message(messages, blocked_msg)}
 
+            if result.status == RailStatus.MODIFIED:
+                log.info("Output modified by rail '%s': content replaced", result.rail or "unknown rail")
+                modified_msg = last_ai_message.model_copy(update={"content": result.content})
+                return {"messages": self._replace_last_ai_message(messages, modified_msg)}
+
             return None
 
         except GuardrailViolation:

tests/integrations/langchain/test_middleware.py

@@ -1342,3 +1342,230 @@ async def test_error_handler_also_replaces_correctly(self, mock_rails_factory):
         assert result["messages"][1].content == middleware.blocked_output_message
         assert isinstance(result["messages"][1], AIMessage)
         assert result["messages"][2] is trailing_msg
+
+
+class TestModifiedStatus:
+    @pytest.mark.asyncio
+    async def test_input_modified_replaces_last_human_message(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="sanitized input")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        state = {"messages": [HumanMessage(content="original input with PII")]}
+        result = await middleware.abefore_model(state, None)
+
+        assert result is not None
+        assert "jump_to" not in result
+        assert len(result["messages"]) == 1
+        assert isinstance(result["messages"][0], HumanMessage)
+        assert result["messages"][0].content == "sanitized input"
+
+    @pytest.mark.asyncio
+    async def test_input_modified_preserves_surrounding_messages(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="redacted")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        system_msg = SystemMessage(content="Be helpful")
+        state = {
+            "messages": [
+                system_msg,
+                HumanMessage(content="my SSN is 123-45-6789"),
+            ]
+        }
+        result = await middleware.abefore_model(state, None)
+
+        assert len(result["messages"]) == 2
+        assert result["messages"][0] is system_msg
+        assert isinstance(result["messages"][1], HumanMessage)
+        assert result["messages"][1].content == "redacted"
+
+    @pytest.mark.asyncio
+    async def test_input_modified_with_multi_turn_history(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="cleaned message")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        first_human = HumanMessage(content="Hello")
+        first_ai = AIMessage(content="Hi there")
+        state = {
+            "messages": [
+                first_human,
+                first_ai,
+                HumanMessage(content="my email is foo@bar.com"),
+            ]
+        }
+        result = await middleware.abefore_model(state, None)
+
+        assert len(result["messages"]) == 3
+        assert result["messages"][0] is first_human
+        assert result["messages"][1] is first_ai
+        assert result["messages"][2].content == "cleaned message"
+
+    @pytest.mark.asyncio
+    async def test_output_modified_replaces_last_ai_message(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="sanitized output")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        state = {
+            "messages": [
+                HumanMessage(content="Hello"),
+                AIMessage(content="original response with PII"),
+            ]
+        }
+        result = await middleware.aafter_model(state, None)
+
+        assert result is not None
+        assert len(result["messages"]) == 2
+        assert isinstance(result["messages"][0], HumanMessage)
+        assert result["messages"][0].content == "Hello"
+        assert isinstance(result["messages"][1], AIMessage)
+        assert result["messages"][1].content == "sanitized output"
+
+    @pytest.mark.asyncio
+    async def test_output_modified_preserves_trailing_messages(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="redacted output")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        trailing = SystemMessage(content="trailing")
+        state = {
+            "messages": [
+                HumanMessage(content="Hello"),
+                AIMessage(content="bad output"),
+                trailing,
+            ]
+        }
+        result = await middleware.aafter_model(state, None)
+
+        assert len(result["messages"]) == 3
+        assert isinstance(result["messages"][1], AIMessage)
+        assert result["messages"][1].content == "redacted output"
+        assert result["messages"][2] is trailing
+
+    @pytest.mark.asyncio
+    async def test_output_modified_replaces_only_last_ai_message(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="fixed")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        state = {
+            "messages": [
+                HumanMessage(content="Hello"),
+                AIMessage(content="First response"),
+                HumanMessage(content="Follow up"),
+                AIMessage(content="Second response with PII"),
+            ]
+        }
+        result = await middleware.aafter_model(state, None)
+
+        assert len(result["messages"]) == 4
+        assert result["messages"][1].content == "First response"
+        assert result["messages"][3].content == "fixed"
+
+    def test_sync_before_model_handles_modified(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="sanitized")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        state = {"messages": [HumanMessage(content="PII content")]}
+        result = middleware.before_model(state, None)
+
+        assert result is not None
+        assert len(result["messages"]) == 1
+        assert isinstance(result["messages"][0], HumanMessage)
+        assert result["messages"][0].content == "sanitized"
+
+    def test_sync_after_model_handles_modified(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="sanitized output")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        state = {
+            "messages": [
+                HumanMessage(content="Hello"),
+                AIMessage(content="PII output"),
+            ]
+        }
+        result = middleware.after_model(state, None)
+
+        assert result is not None
+        assert len(result["messages"]) == 2
+        assert result["messages"][1].content == "sanitized output"
+
+    @pytest.mark.asyncio
+    async def test_input_modified_with_empty_content(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        state = {"messages": [HumanMessage(content="sensitive data")]}
+        result = await middleware.abefore_model(state, None)
+
+        assert result is not None
+        assert len(result["messages"]) == 1
+        assert isinstance(result["messages"][0], HumanMessage)
+        assert result["messages"][0].content == ""
+
+    @pytest.mark.asyncio
+    async def test_input_modified_preserves_message_metadata(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="redacted")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        original = HumanMessage(
+            content="my SSN is 123-45-6789",
+            id="msg-123",
+            name="user1",
+            additional_kwargs={"source": "web"},
+        )
+        state = {"messages": [original]}
+        result = await middleware.abefore_model(state, None)
+
+        modified = result["messages"][0]
+        assert modified.content == "redacted"
+        assert modified.id == "msg-123"
+        assert modified.name == "user1"
+        assert modified.additional_kwargs == {"source": "web"}
+
+    @pytest.mark.asyncio
+    async def test_output_modified_preserves_message_metadata(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="safe output")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        original_ai = AIMessage(
+            content="PII in response",
+            id="ai-456",
+            name="assistant",
+            additional_kwargs={"model": "gpt-4"},
+        )
+        state = {
+            "messages": [
+                HumanMessage(content="Hello"),
+                original_ai,
+            ]
+        }
+        result = await middleware.aafter_model(state, None)
+
+        modified = result["messages"][1]
+        assert modified.content == "safe output"
+        assert modified.id == "ai-456"
+        assert modified.name == "assistant"
+        assert modified.additional_kwargs == {"model": "gpt-4"}
+
+    @pytest.mark.asyncio
+    async def test_output_modified_preserves_tool_calls(self, mock_rails_factory):
+        mock_rails = mock_rails_factory(status=RailStatus.MODIFIED, content="sanitized")
+        middleware = create_middleware_with_rails(mock_rails)
+
+        tool_call = ToolCall(name="search", args={"q": "test"}, id="tc-1")
+        original_ai = AIMessage(
+            content="PII response",
+            id="ai-789",
+            tool_calls=[tool_call],
+        )
+        state = {
+            "messages": [
+                HumanMessage(content="Hello"),
+                original_ai,
+            ]
+        }
+        result = await middleware.aafter_model(state, None)
+
+        modified = result["messages"][1]
+        assert modified.content == "sanitized"
+        assert modified.id == "ai-789"
+        assert len(modified.tool_calls) == 1
+        assert modified.tool_calls[0]["name"] == "search"
+        assert modified.tool_calls[0]["id"] == "tc-1"