fix(memory): redact invalid and legacy tool-call arguments (#6387)

This commit is contained in:
Guo-Yixin 2026-10-06 18:23:56 -05:00 committed by GitHub
parent c5c0dfa54e
commit 789bbaaa0d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 181 additions and 22 deletions

View File

@ -1749,12 +1749,13 @@ call. Skipped or failed compaction leaves the existing messages unchanged.
Optional `pii_redaction.enabled` redacts detected identifiers in user messages,
remote tool results, compaction input, reinjected summaries, and configured
LLM title input. It is off by default. Existing summary placeholders reserve
indices so new values do not reuse them after compaction. No PII mapping is
persisted, so repeated values cannot be linked to a compacted source; numbering
may change when history or summary placeholders disappear. Raw thread text and
local fallback titles remain available for display; memory extraction is outside
this feature's scope.
LLM title input. Memory admission, including pre-compaction flushes, also redacts
detected identifiers in supported text/JSON content, parsed and invalid call arguments/error text,
provider-raw/legacy function calls and supported user-content provenance.
It is off by default; enabled deployments supply a secret for stable, keyed
value-derived placeholders. No token-to-raw mapping is persisted. Raw thread
text and local fallback titles remain available for display; memory redaction
copies messages without changing the caller's history or tool execution.
The Web UI preserves persisted message order when merging history with live updates. Streaming steps around a persisted result inside the loaded history stay together, including steps that arrive after the result. Steps captured during compaction also remain visible before their persisted result when history has not refreshed and the UI has not rendered them yet.

View File

@ -93,6 +93,10 @@ An empty parsed call list does not make an attempted call a final response.
Keep both portable backend filters aligned without importing host helpers;
test the real queue/HTTP write boundary in `tests/test_memory_tool_call_intent.py`.
Memory enqueue redaction also covers `invalid_tool_calls` arguments/error text
and legacy `function_call` payloads. Keep sync, async and compaction admission
aligned; preserve original messages and detector policy.
`memory.mode: tool` registers the four memory tools.
The model chooses when to search or change facts.
Tool mode still uses `MemoryMiddleware` for passive writes on supported remote backends.

View File

@ -63,7 +63,7 @@ No match: skip inspection, log ID. Tests: `test_read_mark_tool_call_correlation.
ordered by application — the last entry produced the final visible bytes — so an
observer classifies raw→visible transforms from facts rather than by sniffing
output wording.
4. **PiiRedactionMiddleware** - *(optional, `pii_redaction.enabled`, default off, #3190)* Rewrites PII in genuine user messages (`wrap_model_call`) and remote-content tool results (`wrap_tool_call`, the tool-result allowlist) to irreversible value-derived placeholders. Deterministic regex detectors only. Innermost Layer-1 wrapper; compaction input, summaries, title input and queued memory payloads are redacted via `redact_text`.
4. **PiiRedactionMiddleware** - *(optional, `pii_redaction.enabled`, default off, #3190)* Rewrites PII in genuine user messages (`wrap_model_call`) and remote-content tool results (`wrap_tool_call`, the tool-result allowlist) to irreversible value-derived placeholders. Deterministic regex detectors only. Innermost Layer-1 wrapper; compaction input, summaries and title input use `redact_text`. Memory admission redacts content, parsed/invalid/legacy call arguments and supported provenance before backend filtering.
5. **ThreadDataMiddleware** - Creates per-thread directories under the user's isolation scope (`backend/.deer-flow/users/{user_id}/threads/{thread_id}/user-data/{workspace,uploads,outputs}`); resolves identity via `resolve_runtime_user_id(runtime)`, including Gateway runtime context and standalone LangGraph Server auth, then falls back to the request ContextVar / `"default"`
6. **UploadsMiddleware** - Injects current uploads (lead only). Check existence in the runtime user's bucket; coerce invalid `files[*].size` to `0`. Copy the user message when adding context to preserve caller metadata, including `response_metadata`.

View File

@ -33,8 +33,8 @@ def redact_queued_messages(messages: list, pii_redaction_config: PiiRedactionCon
renders the same token across the batch, across enqueues, and across
seams; message objects are rebuilt rather than mutated. Structured
tool-call arguments are covered too — backends like OpenViking retain the full
message object, including ``tool_calls`` and provider-format arguments in
``additional_kwargs``. Text-bearing provenance metadata is covered as
message object, including parsed/invalid calls (and invalid-call errors)
and provider-raw/legacy arguments in ``additional_kwargs``. Text-bearing provenance metadata is covered as
well: UploadsMiddleware preserves the raw user turn in
``original_user_content``, and clarification replies keep the raw answer
in ``human_input_response.value`` — DeerMem reads that mapping — so a
@ -49,18 +49,20 @@ def redact_queued_messages(messages: list, pii_redaction_config: PiiRedactionCon
content, changed = _redact_content(message.content, redactor)
if changed:
updates["content"] = content
tool_calls = getattr(message, "tool_calls", None)
if tool_calls:
rewritten = [_redact_tool_call(call, redactor) for call in tool_calls]
if rewritten != tool_calls:
updates["tool_calls"] = rewritten
for field in ("tool_calls", "invalid_tool_calls"):
tool_calls = getattr(message, field, None)
if tool_calls:
rewritten = [_redact_tool_call(call, redactor) for call in tool_calls]
if rewritten != tool_calls:
updates[field] = rewritten
additional_kwargs = getattr(message, "additional_kwargs", None)
if isinstance(additional_kwargs, dict):
new_additional_kwargs = additional_kwargs
if additional_kwargs.get("tool_calls"):
rewritten = _redact_strings(additional_kwargs["tool_calls"], redactor)
if rewritten != additional_kwargs["tool_calls"]:
new_additional_kwargs = {**new_additional_kwargs, "tool_calls": rewritten}
for field in ("tool_calls", "function_call"):
if additional_kwargs.get(field):
rewritten = _redact_strings(additional_kwargs[field], redactor)
if rewritten != additional_kwargs[field]:
new_additional_kwargs = {**new_additional_kwargs, field: rewritten}
original_content = additional_kwargs.get(ORIGINAL_USER_CONTENT_KEY)
if isinstance(original_content, str) and original_content:
redacted_original = redactor.redact(original_content)
@ -78,11 +80,12 @@ def redact_queued_messages(messages: list, pii_redaction_config: PiiRedactionCon
def _redact_tool_call(call: dict, redactor: _Redactor) -> dict:
"""Redact PII inside one parsed tool-call dict (args values recursively)."""
"""Redact arguments and invalid-call error text without changing call identity."""
rewritten = dict(call)
args = rewritten.get("args")
if args is not None:
rewritten["args"] = _redact_strings(args, redactor)
for field in ("args", "error"):
value = rewritten.get(field)
if value is not None:
rewritten[field] = _redact_strings(value, redactor)
return rewritten

View File

@ -0,0 +1,102 @@
"""PII redaction must cover malformed and legacy calls at memory admission."""
from __future__ import annotations
import asyncio
from unittest.mock import AsyncMock, MagicMock
import pytest
from langchain_core.messages import AIMessage, HumanMessage, messages_from_dict, messages_to_dict
from langchain_openai.chat_models.base import _convert_dict_to_message
from langgraph.runtime import Runtime
from deerflow.agents.memory import summarization_hook
from deerflow.agents.middlewares import memory_middleware
from deerflow.agents.middlewares.memory_middleware import MemoryMiddleware
from deerflow.agents.middlewares.pii_redaction_middleware import redact_text
from deerflow.agents.middlewares.summarization_middleware import SummarizationEvent
from deerflow.config.memory_config import MemoryConfig
from deerflow.config.pii_redaction_config import PiiRedactionConfig
_EMAIL = "alice@example.com"
_SECRET = "memory-regression-deployment-secret"
def _attempt(kind: str) -> AIMessage:
if kind == "invalid-checkpoint":
# A persisted parsed-invalid view need not retain provider-raw calls.
message = AIMessage(content="Preparing contact lookup", invalid_tool_calls=[{"name": "lookup", "id": "call-1", "args": '{"email":"alice@example.com",', "error": "Invalid arguments for alice@example.com"}])
elif kind == "legacy":
message = _convert_dict_to_message({"role": "assistant", "content": "Preparing contact lookup", "function_call": {"name": "lookup", "arguments": '{"email":"alice@example.com"}'}})
else:
message = _convert_dict_to_message({"role": "assistant", "content": "Preparing contact lookup", "tool_calls": [{"id": "call-1", "type": "function", "function": {"name": "lookup", "arguments": '{"email":"alice@example.com",'}}]})
assert isinstance(message, AIMessage) and not message.tool_calls and message.invalid_tool_calls
assert isinstance(message, AIMessage)
return messages_from_dict(messages_to_dict([message]))[0]
def _capture(monkeypatch: pytest.MonkeyPatch, entry: str, messages: list, pii_config: PiiRedactionConfig | None) -> list:
recorded = []
def record(_thread_id, queued, **_kwargs):
recorded.append(queued)
manager = MagicMock()
manager.add.side_effect = record
manager.aadd = AsyncMock(side_effect=record)
manager.add_nowait.side_effect = record
monkeypatch.setattr(memory_middleware, "get_memory_manager", lambda: manager)
monkeypatch.setattr(summarization_hook, "get_memory_manager", lambda: manager)
monkeypatch.setattr(summarization_hook, "get_memory_config", lambda: MemoryConfig(enabled=True))
runtime = Runtime(context={"thread_id": "thread-1", "user_id": "alice"})
middleware = MemoryMiddleware(memory_config=MemoryConfig(enabled=True), pii_redaction_config=pii_config)
if entry == "after-agent":
middleware.after_agent({"messages": messages}, runtime)
elif entry == "async-after-agent":
asyncio.run(middleware.aafter_agent({"messages": messages}, runtime))
else:
event = SummarizationEvent(messages_to_summarize=tuple(messages), preserved_messages=(), thread_id="thread-1", agent_name=None, runtime=runtime)
summarization_hook.memory_flush_hook(event, pii_redaction_config=pii_config)
assert len(recorded) == 1
return recorded[0]
@pytest.mark.parametrize("kind", ["provider-invalid", "invalid-checkpoint", "legacy"])
@pytest.mark.parametrize("entry", ["after-agent", "async-after-agent", "compaction"])
def test_invalid_and_legacy_call_views_are_redacted_before_memory_admission(monkeypatch, kind: str, entry: str) -> None:
attempt = _attempt(kind)
messages = [HumanMessage(f"Look up {_EMAIL}"), attempt, AIMessage("Lookup was not executed")]
original = messages_to_dict(messages)
config = PiiRedactionConfig(enabled=True, token_secret=_SECRET)
queued = _capture(monkeypatch, entry, messages, config)
assert _EMAIL not in str(messages_to_dict(queued))
token = redact_text(_EMAIL, config)
assert token in queued[0].content
if kind == "legacy":
assert queued[1].additional_kwargs["function_call"]["name"] == "lookup"
assert token in queued[1].additional_kwargs["function_call"]["arguments"]
else:
assert not queued[1].tool_calls
assert queued[1].invalid_tool_calls[0]["id"] == "call-1"
assert queued[1].invalid_tool_calls[0]["name"] == "lookup"
assert token in queued[1].invalid_tool_calls[0]["args"]
assert token in queued[1].invalid_tool_calls[0]["error"]
assert queued[1].content == attempt.content
assert queued[2] is messages[2]
assert messages_to_dict(messages) == original
@pytest.mark.parametrize("kind", ["provider-invalid", "invalid-checkpoint", "legacy"])
@pytest.mark.parametrize("entry", ["after-agent", "async-after-agent", "compaction"])
@pytest.mark.parametrize("policy", ["no-config", "disabled", "email-disabled"])
def test_call_views_follow_existing_redaction_opt_in_and_detector_policy(monkeypatch, kind: str, entry: str, policy: str) -> None:
config = None if policy == "no-config" else PiiRedactionConfig(enabled=policy != "disabled", redact_email=policy != "email-disabled", token_secret=_SECRET)
messages = [_attempt(kind)]
original = messages_to_dict(messages)
queued = _capture(monkeypatch, entry, messages, config)
assert queued[0] is messages[0]
assert messages_to_dict(queued) == original

View File

@ -163,6 +163,55 @@ _ACTOR_PEER: ContextVar[str | None] = ContextVar(
)
@pytest.mark.parametrize("kind", ["invalid", "legacy"])
@pytest.mark.parametrize("entry", ["after-agent", "async-after-agent", "compaction"])
def test_memory_admission_redacts_call_views_before_openviking_recorder(tmp_path, monkeypatch, official_integration, kind: str, entry: str) -> None:
import asyncio
from langchain_core.messages import messages_to_dict
from langchain_openai.chat_models.base import _convert_dict_to_message
from langgraph.runtime import Runtime
from deerflow.agents.memory import summarization_hook
from deerflow.agents.middlewares import memory_middleware
from deerflow.agents.middlewares.memory_middleware import MemoryMiddleware
from deerflow.agents.middlewares.summarization_middleware import SummarizationEvent
from deerflow.config.memory_config import MemoryConfig
from deerflow.config.pii_redaction_config import PiiRedactionConfig
config = PiiRedactionConfig(enabled=True, token_secret="openviking-regression-secret")
raw = {"role": "assistant", "content": "Preparing lookup"}
call = {"name": "lookup", "arguments": '{"email":"alice@example.com"}' if kind == "legacy" else '{"email":"alice@example.com",'}
raw.update({"function_call": call} if kind == "legacy" else {"tool_calls": [{"id": "call-1", "type": "function", "function": call}]})
messages = [HumanMessage("Look up contact"), _convert_dict_to_message(raw), AIMessage("Lookup was not executed")]
original = messages_to_dict(messages)
manager = _manager(tmp_path, monkeypatch)
monkeypatch.setattr(memory_middleware, "get_memory_manager", lambda: manager)
monkeypatch.setattr(summarization_hook, "get_memory_manager", lambda: manager)
monkeypatch.setattr(summarization_hook, "get_memory_config", lambda: MemoryConfig(enabled=True))
runtime = Runtime(context={"thread_id": "thread-1", "user_id": "alice"})
middleware = MemoryMiddleware(memory_config=MemoryConfig(enabled=True), pii_redaction_config=config)
try:
if entry == "after-agent":
middleware.after_agent({"messages": messages}, runtime)
elif entry == "async-after-agent":
asyncio.run(middleware.aafter_agent({"messages": messages}, runtime))
else:
event = SummarizationEvent(messages_to_summarize=tuple(messages), preserved_messages=(), thread_id="thread-1", agent_name=None, runtime=runtime)
summarization_hook.memory_flush_hook(event, pii_redaction_config=config)
assert len(manager._recorder.calls) == 1
recorded = manager._recorder.calls[0][1]
assert len(recorded) == 3
assert "alice@example.com" not in str(messages_to_dict(recorded))
if kind == "legacy":
assert "[EMAIL_" in recorded[1].additional_kwargs["function_call"]["arguments"]
else:
assert "[EMAIL_" in recorded[1].invalid_tool_calls[0]["args"]
assert messages_to_dict(messages) == original
finally:
manager.close()
@contextmanager
def _use_actor_peer(peer_id: str | None):
token = _ACTOR_PEER.set(peer_id)