From 789bbaaa0da37da9653cae9833b3846fde484fc4 Mon Sep 17 00:00:00 2001 From: Guo-Yixin Date: Tue, 6 Oct 2026 18:23:56 -0500 Subject: [PATCH] fix(memory): redact invalid and legacy tool-call arguments (#6387) --- README.md | 13 +-- .../harness/deerflow/agents/memory/AGENTS.md | 4 + .../deerflow/agents/middlewares/AGENTS.md | 2 +- .../agents/middlewares/memory_middleware.py | 33 +++--- .../tests/test_memory_tool_call_redaction.py | 102 ++++++++++++++++++ .../tests/test_openviking_memory_backend.py | 49 +++++++++ 6 files changed, 181 insertions(+), 22 deletions(-) create mode 100644 backend/tests/test_memory_tool_call_redaction.py diff --git a/README.md b/README.md index d3cc73922d..b87d9a3f0f 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/backend/packages/harness/deerflow/agents/memory/AGENTS.md b/backend/packages/harness/deerflow/agents/memory/AGENTS.md index e7e8f40d16..363df137bc 100644 --- a/backend/packages/harness/deerflow/agents/memory/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/memory/AGENTS.md @@ -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. diff --git a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md index 4325b70b6d..b8297a7644 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md @@ -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`. diff --git a/backend/packages/harness/deerflow/agents/middlewares/memory_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/memory_middleware.py index 69e02b734b..b9cc2ce6b5 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/memory_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/memory_middleware.py @@ -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 diff --git a/backend/tests/test_memory_tool_call_redaction.py b/backend/tests/test_memory_tool_call_redaction.py new file mode 100644 index 0000000000..a1b4c70bda --- /dev/null +++ b/backend/tests/test_memory_tool_call_redaction.py @@ -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 diff --git a/backend/tests/test_openviking_memory_backend.py b/backend/tests/test_openviking_memory_backend.py index 267e65a569..3a7547cb52 100644 --- a/backend/tests/test_openviking_memory_backend.py +++ b/backend/tests/test_openviking_memory_backend.py @@ -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)