langchain - 💡(How to fix) Fix PIIMiddleware state hooks skip tool-call-only AI messages and stringify structured content

Checked other resources

• This is a bug, not a usage question.
• I added a clear and descriptive title that summarizes the issue.
• I used the GitHub search to find similar issues and did not find this exact state-hook regression.
• I am sure this is a bug in LangChain rather than in user code.
• This is not related to langchain-community.

Package (Required)

• langchain
• langchain-openai
• langchain-anthropic
• langchain-classic
• langchain-core
• langchain-model-profiles
• langchain-tests
• langchain-text-splitters
• Other / not sure / general

Related Issues / PRs

• Related background: #35011
• Original PIIMiddleware implementation: #33271
• Recent streaming/value-snapshot redaction work: #37616

Description

PIIMiddleware now has two redaction paths:

1. state-level hooks (before_model / after_model)
2. the recursive message/value walker added in #37616 for streaming surfaces and values snapshots

Those two paths disagree on two important cases:

1. after_model returns early for tool-call-only AIMessage(content="", tool_calls=[...]), so tool_calls[*].args are not redacted for non-streaming .invoke() / canonical graph state.
2. the state hooks call str(message.content) for HumanMessage, ToolMessage, and AIMessage, which stringifies list-typed structured content blocks instead of preserving the list shape.

This looks like historical drift rather than intended behavior:

• the original middleware in #33271 implemented the state hooks with manual str(message.content) handling
• #37616 added _redact_value() / _redact_base_message(), added unit coverage for empty-content tool-call messages and structured content, and documented the state-level hooks as the backstop for non-streaming consumers
• the state hooks were not switched over to the shared helper, so the canonical state path still behaves differently from the newer recursive redactor

Reproduction Steps / Example Code (Python)

from typing import Any

from langchain.agents import AgentState
from langchain.agents.middleware.pii import PIIMiddleware
from langchain_core.messages import AIMessage, HumanMessage, ToolCall
from langgraph.runtime import Runtime

# 1) tool-call-only AI messages skip output redaction
middleware = PIIMiddleware("email", apply_to_output=True)
msg = AIMessage(
    content="",
    tool_calls=[ToolCall(name="send_email", args={"to": "[email protected]"}, id="c1")],
)
state = AgentState[Any](messages=[HumanMessage("hi"), msg])
result = middleware.after_model(state, Runtime())
assert result is not None
assert result["messages"][-1].tool_calls[0]["args"] == {"to": "[REDACTED_EMAIL]"}

# 2) structured content is stringified instead of preserved
middleware = PIIMiddleware("email", apply_to_input=True)
msg = HumanMessage(content=[{"type": "text", "text": "Reach me at [email protected]"}])
state = AgentState[Any](messages=[msg])
result = middleware.before_model(state, Runtime())
assert result is not None
assert isinstance(result["messages"][0].content, list)
assert result["messages"][0].content[0]["text"] == "Reach me at [REDACTED_EMAIL]"

Actual behavior

• after_model returns None for empty-content tool-call messages because it gates on not last_ai_msg.content
• the state hooks use str(message.content), turning structured content into a Python repr string like:

"[{'type': 'text', 'text': 'Reach me at [REDACTED_EMAIL]'}]"

Expected behavior

State-level hooks should redact the same message surfaces the shared recursive helper already handles:

• HumanMessage.content and ToolMessage.content when they are list-typed content blocks
• AIMessage.content
• AIMessage.tool_calls[*].args
• AIMessage.invalid_tool_calls[*].args

In particular, after_model should not skip tool-call-only assistant messages just because content == "".

Why this looks unintentional

• _redact_value() / _redact_base_message() already cover both cases and already have unit coverage in test_pii.py
• the PIIMiddleware docstring added in #37616 says state-level hooks remain the backstop for non-streaming consumers
• today the streaming/values path and the canonical state path disagree on exactly those cases

Suggested fix

Refactor before_model and after_model to reuse the shared recursive message redactor instead of manually calling str(message.content) and reconstructing partial message objects.

That should also preserve message shape and metadata more reliably than the current manual reconstruction.

System Info

Observed on langchain-ai/langchain master at 33875fde2acf6ffb717915a895638274a6098ec2 (May 22, 2026), after #37616 was merged on May 22, 2026.