From 4a66e355da07b8fad7d6cb8db99287fb4d7e6d83 Mon Sep 17 00:00:00 2001 From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:46:14 -0700 Subject: [PATCH] fix(langchain): preserve content-block shape in PIIMiddleware redaction (#39894) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `PIIMiddleware` redacted message content via `str(message.content)`. When `.content` is a list of content blocks rather than a plain string, `str(...)` produces the `repr` of the list — so a redacted user message was stored as the literal string: ``` "[{'type': 'text', 'text': 'my email is [REDACTED_EMAIL]'}]" ``` The PII itself was still detected and redacted; only the shape of the message was destroyed. Anyone sending multimodal or provider-native block content through an agent with `PIIMiddleware` had their messages silently flattened into that repr before reaching the model. `_process_content` now accepts either shape and walks content blocks, redacting the text of each block in place and leaving non-text blocks untouched. Plain-string content behaves exactly as before. While fixing the call sites, this also drops the hand-rebuilding of messages (`HumanMessage(content=..., id=..., name=...)`) in favor of `model_copy`. The enumerated rebuild had been quietly discarding every field it did not list — `additional_kwargs`, `response_metadata`, and so on — on any message that contained PII. ## Release note Fixed `PIIMiddleware` flattening list-of-content-blocks message content into its string `repr` during redaction. Redaction now preserves the original content shape, and no longer drops message fields such as `additional_kwargs` on redacted messages. --- *Written with the help of an AI agent (Claude Code).* --- .../langchain/agents/middleware/pii.py | 88 +++++++++---------- .../middleware/implementations/test_pii.py | 83 +++++++++++++++++ 2 files changed, 127 insertions(+), 44 deletions(-) diff --git a/libs/langchain_v1/langchain/agents/middleware/pii.py b/libs/langchain_v1/langchain/agents/middleware/pii.py index c3b99eee8a..e4c7612731 100644 --- a/libs/langchain_v1/langchain/agents/middleware/pii.py +++ b/libs/langchain_v1/langchain/agents/middleware/pii.py @@ -5,7 +5,7 @@ from __future__ import annotations from functools import partial from typing import TYPE_CHECKING, Any, ClassVar, Literal -from langchain_core.messages import AIMessage, AnyMessage, BaseMessage, HumanMessage, ToolMessage +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage from langgraph.stream import StreamTransformer from typing_extensions import override @@ -660,13 +660,39 @@ class PIIMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT]) """Name of the middleware.""" return f"{self.__class__.__name__}[{self.pii_type}]" - def _process_content(self, content: str) -> tuple[str, list[PIIMatch]]: - """Apply the configured redaction rule to the provided content.""" - matches = self.detector(content) - if not matches: - return content, [] - sanitized = apply_strategy(content, matches, self.strategy) - return sanitized, matches + def _process_content( + self, content: str | list[str | dict[str, Any]] + ) -> tuple[str | list[str | dict[str, Any]], bool]: + """Apply the configured redaction rule, preserving the shape of `content`. + + Only string leaves are redacted: a plain-string `content` directly, and + the `text` of each block in a block-list `content`. Non-text blocks pass + through untouched, so block-list content keeps its structure instead of + collapsing into the `repr` of the list. + + Args: + content: The message content to redact. + + Returns: + The redacted content and whether anything was rewritten. + """ + if isinstance(content, str): + matches = self.detector(content) + if not matches: + return content, False + return apply_strategy(content, matches, self.strategy), True + + new_blocks: list[str | dict[str, Any]] = [] + changed = False + for block in content: + new_block: str | dict[str, Any] = block + text = block if isinstance(block, str) else block.get("text") + if isinstance(text, str) and (matches := self.detector(text)): + redacted = apply_strategy(text, matches, self.strategy) + new_block = redacted if isinstance(block, str) else {**block, "text": redacted} + changed = True + new_blocks.append(new_block) + return (new_blocks, True) if changed else (content, False) @hook_config(can_jump_to=["end"]) @override @@ -710,18 +736,12 @@ class PIIMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT]) break if last_user_idx is not None and last_user_msg and last_user_msg.content: - # Detect PII in message content - content = str(last_user_msg.content) - new_content, matches = self._process_content(content) + new_content, changed = self._process_content(last_user_msg.content) - if matches: - updated_message: AnyMessage = HumanMessage( - content=new_content, - id=last_user_msg.id, - name=last_user_msg.name, + if changed: + new_messages[last_user_idx] = last_user_msg.model_copy( + update={"content": new_content} ) - - new_messages[last_user_idx] = updated_message any_modified = True # Check tool results if enabled @@ -742,21 +762,12 @@ class PIIMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT]) if not tool_msg.content: continue - content = str(tool_msg.content) - new_content, matches = self._process_content(content) + new_content, changed = self._process_content(tool_msg.content) - if not matches: + if not changed: continue - # Create updated tool message - updated_message = ToolMessage( - content=new_content, - id=tool_msg.id, - name=tool_msg.name, - tool_call_id=tool_msg.tool_call_id, - ) - - new_messages[i] = updated_message + new_messages[i] = tool_msg.model_copy(update={"content": new_content}) any_modified = True if any_modified: @@ -824,24 +835,13 @@ class PIIMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT]) if last_ai_idx is None or not last_ai_msg or not last_ai_msg.content: return None - # Detect PII in message content - content = str(last_ai_msg.content) - new_content, matches = self._process_content(content) + new_content, changed = self._process_content(last_ai_msg.content) - if not matches: + if not changed: return None - # Create updated message - updated_message = AIMessage( - content=new_content, - id=last_ai_msg.id, - name=last_ai_msg.name, - tool_calls=last_ai_msg.tool_calls, - ) - - # Return updated messages new_messages = list(messages) - new_messages[last_ai_idx] = updated_message + new_messages[last_ai_idx] = last_ai_msg.model_copy(update={"content": new_content}) return {"messages": new_messages} diff --git a/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_pii.py b/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_pii.py index 4945e49a5f..a86ed93164 100644 --- a/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_pii.py +++ b/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_pii.py @@ -551,6 +551,89 @@ class TestPIIMiddlewareIntegration: messages = result["messages"] assert any("[REDACTED_EMAIL]" in str(msg.content) for msg in messages) + def test_input_list_content_preserved(self) -> None: + """List-of-content-blocks input is redacted in place, not stringified. + + Regression: `str(msg.content)` flattened block-list content into its + `repr`, so the stored content became the literal string + `"[{'type': 'text', 'text': '...'}]"`. + """ + middleware = PIIMiddleware("email", strategy="redact") + state = AgentState[Any]( + messages=[ + HumanMessage(content=[{"type": "text", "text": "my email is test@example.com"}]) + ] + ) + + result = middleware.before_model(state, Runtime()) + + assert result is not None + content = result["messages"][0].content + assert content == [{"type": "text", "text": "my email is [REDACTED_EMAIL]"}] + assert "test@example.com" not in str(content) + + def test_output_list_content_preserved(self) -> None: + """List-of-content-blocks AI output is redacted in place, not stringified.""" + middleware = PIIMiddleware( + "email", strategy="redact", apply_to_input=False, apply_to_output=True + ) + state = AgentState[Any]( + messages=[AIMessage(content=[{"type": "text", "text": "reach me at ai@example.com"}])] + ) + + result = middleware.after_model(state, Runtime()) + + assert result is not None + content = result["messages"][0].content + assert content == [{"type": "text", "text": "reach me at [REDACTED_EMAIL]"}] + assert "ai@example.com" not in str(content) + + def test_redaction_preserves_message_fields(self) -> None: + """Redacting keeps fields the hooks never enumerated when rebuilding messages.""" + middleware = PIIMiddleware("email", strategy="redact") + original = HumanMessage( + content="my email is test@example.com", + id="msg_1", + name="alice", + additional_kwargs={"source": "web"}, + ) + state = AgentState[Any](messages=[original]) + + result = middleware.before_model(state, Runtime()) + + assert result is not None + redacted = result["messages"][0] + assert redacted.content == "my email is [REDACTED_EMAIL]" + assert redacted.id == "msg_1" + assert redacted.name == "alice" + assert redacted.additional_kwargs == {"source": "web"} + + def test_tool_result_list_content_preserved(self) -> None: + """List-of-content-blocks tool results are redacted in place.""" + middleware = PIIMiddleware( + "email", strategy="redact", apply_to_input=False, apply_to_tool_results=True + ) + state = AgentState[Any]( + messages=[ + HumanMessage("Search for user"), + AIMessage( + content="", + tool_calls=[ToolCall(name="search", args={}, id="call_1", type="tool_call")], + ), + ToolMessage( + content=[{"type": "text", "text": "found: john@example.com"}], + tool_call_id="call_1", + ), + ] + ) + + result = middleware.before_model(state, Runtime()) + + assert result is not None + content = result["messages"][2].content + assert content == [{"type": "text", "text": "found: [REDACTED_EMAIL]"}] + assert "john@example.com" not in str(content) + class TestCustomDetector: """Test custom detector functionality."""