From 7e64330dd1f05b0e9f0c817f124058f6170966ee Mon Sep 17 00:00:00 2001 From: lilei <93921041+xihongshichaojidan8@users.noreply.github.com> Date: Fri, 6 Nov 2026 21:24:57 +0800 Subject: [PATCH] fix(memory): avoid recursive tool-argument admission (#7928) * fix(memory): avoid recursive tool-argument admission * fix(memory): track argument redaction changes without deep equality --- README.md | 3 + .../deerflow/agents/middlewares/AGENTS.md | 2 +- .../agents/middlewares/memory_middleware.py | 78 +++++-- backend/tests/test_memory_admission_depth.py | 212 ++++++++++++++++++ 4 files changed, 274 insertions(+), 21 deletions(-) create mode 100644 backend/tests/test_memory_admission_depth.py diff --git a/README.md b/README.md index 0bcf7b927b..615ea74f28 100644 --- a/README.md +++ b/README.md @@ -2721,6 +2721,9 @@ 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. +For SDK agents without a checkpointer, deeply nested JSON tool arguments do not +prevent memory admission, with PII redaction either enabled or disabled. +Structured JSON content blocks still use recursive redaction and remain depth-limited. The hidden `/goal` continuation is redacted before it is stored. 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/middlewares/AGENTS.md b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md index 08e1835875..30ac852bde 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md @@ -98,7 +98,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 and title input use `redact_text`. Memory admission redacts content, parsed/invalid/legacy call arguments and supported provenance before backend filtering. +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. Disabled PII skips conversion; tool arguments track changes without deep equality. JSON content blocks remain recursive. SDK tests: `tests/test_memory_admission_depth.py`. 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** - Lead-only current uploads: exact basenames in runtime owner's bucket; skip malformed references/probe errors (debug-log errors). Invalid/display-overflowing sizes use `0`; render >=1 TB compactly without changing stored sizes. Copy messages with 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 608d27d7cb..074d7019d4 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/memory_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/memory_middleware.py @@ -32,7 +32,8 @@ def redact_queued_messages(messages: list, pii_redaction_config: PiiRedactionCon Shared by MemoryMiddleware's enqueue boundary and the compaction-triggered ``memory_flush_hook``: placeholders are value-derived, so the same identity renders the same token across the batch, across enqueues, and across - seams; message objects are rebuilt rather than mutated. Structured + seams. Changed messages are copied; unchanged/disabled messages are + shared with the caller. Queue consumers must not mutate them. Structured tool-call arguments are covered too — backends like OpenViking retain the full 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 @@ -42,6 +43,9 @@ def redact_queued_messages(messages: list, pii_redaction_config: PiiRedactionCon redacted ``content`` alone would still leak either one to the memory backend. """ + if pii_redaction_config is None or not pii_redaction_config.enabled: + return list(messages) + redactor = _make_redactor(pii_redaction_config) redacted = [] @@ -53,16 +57,21 @@ def redact_queued_messages(messages: list, pii_redaction_config: PiiRedactionCon 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: + rewritten = [] + calls_changed = False + for call in tool_calls: + rewritten_call, call_changed = _redact_tool_call(call, redactor) + rewritten.append(rewritten_call) + calls_changed = calls_changed or call_changed + if calls_changed: updates[field] = rewritten additional_kwargs = getattr(message, "additional_kwargs", None) if isinstance(additional_kwargs, dict): new_additional_kwargs = additional_kwargs for field in ("tool_calls", "function_call"): if additional_kwargs.get(field): - rewritten = _redact_strings(additional_kwargs[field], redactor) - if rewritten != additional_kwargs[field]: + rewritten, arguments_changed = _redact_strings(additional_kwargs[field], redactor) + if arguments_changed: 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: @@ -80,18 +89,20 @@ def redact_queued_messages(messages: list, pii_redaction_config: PiiRedactionCon return redacted -def _redact_tool_call(call: dict, redactor: _Redactor) -> dict: +def _redact_tool_call(call: dict, redactor: _Redactor) -> tuple[dict, bool]: """Redact arguments and invalid-call error text without changing call identity.""" rewritten = dict(call) + changed = False for field in ("args", "error"): value = rewritten.get(field) if value is not None: - rewritten[field] = _redact_strings(value, redactor) - return rewritten + rewritten[field], field_changed = _redact_strings(value, redactor) + changed = changed or field_changed + return rewritten, changed -def _redact_strings(value: object, redactor: _Redactor) -> object: - """Redact every string leaf in a JSON-like *value* tree, keys included. +def _redact_strings(value: object, redactor: _Redactor) -> tuple[object, bool]: + """Redact string leaves/keys and track changes without comparing trees. Tool arguments can carry user data in mapping keys (e.g. ``{"contacts": {"alice@example.com": "manager"}}``), so string keys are @@ -101,16 +112,43 @@ def _redact_strings(value: object, redactor: _Redactor) -> object: accepted, since at 128 bits that only happens for identical identities. """ if isinstance(value, str): - return redactor.redact(value) - if isinstance(value, list): - return [_redact_strings(item, redactor) for item in value] - if isinstance(value, dict): - rebuilt = {} - for key, item in value.items(): - new_key = redactor.redact(key) if isinstance(key, str) else key - rebuilt[new_key] = _redact_strings(item, redactor) - return rebuilt - return value + redacted = redactor.redact(value) + return redacted, redacted != value + if not isinstance(value, (dict, list)): + return value, False + + changed = False + rebuilt = {} if isinstance(value, dict) else [] + stack = [(value, iter(value.items()) if isinstance(value, dict) else enumerate(value), rebuilt)] + ancestors = {id(value)} + while stack: + source, items, target = stack[-1] + try: + key, item = next(items) + except StopIteration: + ancestors.remove(id(source)) + stack.pop() + continue + + new_key = redactor.redact(key) if isinstance(key, str) else key + if isinstance(key, str): + changed = changed or new_key != key + if isinstance(item, str): + rewritten = redactor.redact(item) + changed = changed or rewritten != item + elif isinstance(item, (dict, list)): + if id(item) in ancestors: + raise ValueError("Cyclic tool arguments are not supported JSON") + rewritten = {} if isinstance(item, dict) else [] + ancestors.add(id(item)) + stack.append((item, iter(item.items()) if isinstance(item, dict) else enumerate(item), rewritten)) + else: + rewritten = item + if isinstance(target, dict): + target[new_key] = rewritten + else: + target.append(rewritten) + return rebuilt, changed class MemoryMiddlewareState(AgentState): diff --git a/backend/tests/test_memory_admission_depth.py b/backend/tests/test_memory_admission_depth.py new file mode 100644 index 0000000000..57daa8ef13 --- /dev/null +++ b/backend/tests/test_memory_admission_depth.py @@ -0,0 +1,212 @@ +"""Memory admission must not abort SDK graphs on valid deep tool arguments.""" + +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage +from langchain_core.tools import tool + +from deerflow.agents.factory import create_deerflow_agent +from deerflow.agents.features import RuntimeFeatures +from deerflow.agents.middlewares import memory_middleware +from deerflow.agents.middlewares.pii_redaction_middleware import redact_text +from deerflow.config.memory_config import MemoryConfig +from deerflow.config.pii_redaction_config import PiiRedactionConfig + +_EMAIL = "alice@example.com" +_OTHER_EMAIL = "bob@example.com" + + +class _ScriptedModel(FakeMessagesListChatModel): + def bind_tools(self, tools, **kwargs): + return self + + +@tool +def _consume(nested: dict) -> str: + """Consume a nested JSON structure.""" + return "done" + + +def _arguments(depth: int, shape: str = "dict") -> dict: + value = {_EMAIL: _OTHER_EMAIL, "number": 1, "flag": False, "empty": None} + for _ in range(depth): + value = {"nested": value} if shape == "dict" else {"nested": [value]} + # Keep this a valid-JSON regression, not an arbitrary Python-only state. + return json.loads(json.dumps(value)) + + +def _json_depth_limit(maximum: int, shape: str = "dict") -> int: + """Find an accepted JSON depth without changing interpreter recursion limits.""" + low, high = 0, maximum + while low < high: + middle = (low + high + 1) // 2 + try: + _arguments(middle, shape) + except RecursionError: + high = middle - 1 + else: + low = middle + return low + + +@pytest.fixture(params=["shallow", "json-boundary"]) +def depth(request): + # A fixed 1100 may exceed JSON limits and skip on some interpreters. + # Probe the accepted boundary so every deep case actually runs. + return request.param + + +def _leaf(value: dict, depth: int, shape: str = "dict") -> dict: + for _ in range(depth): + value = value["nested"] + if shape == "mixed": + value = value[0] + return value + + +def _policy(policy: str) -> PiiRedactionConfig | None: + return None if policy == "omitted" else PiiRedactionConfig(enabled=policy != "disabled", redact_email=policy != "email-disabled", token_secret="memory-depth-regression-secret") + + +@pytest.mark.parametrize("mode", ["sync", "async"]) +@pytest.mark.parametrize("memory", [False, True]) +@pytest.mark.parametrize("policy", ["omitted", "disabled", "enabled", "email-disabled"]) +@pytest.mark.asyncio +async def test_sdk_completion_with_deep_memory_arguments(monkeypatch, mode, depth, memory, policy): + depth = 10 if depth == "shallow" else _json_depth_limit(10_000) + args = _arguments(depth) + original = json.dumps(args) + attempt = AIMessage(content="", id="historical-ai", tool_calls=[{"name": "_consume", "args": args, "id": "old-call"}]) + messages = [HumanMessage("Analyze."), attempt, ToolMessage(content="done", tool_call_id="old-call"), HumanMessage("Summarize.")] + manager = SimpleNamespace(add=Mock(), aadd=AsyncMock()) + monkeypatch.setattr(memory_middleware, "get_memory_manager", lambda: manager) + config = _policy(policy) + graph = create_deerflow_agent( + _ScriptedModel(responses=[AIMessage(content="Task complete.")]), + tools=[_consume], + features=RuntimeFeatures(sandbox=False, memory=memory, memory_config=MemoryConfig(enabled=True), loop_detection=False), + checkpointer=None, + **({} if policy == "omitted" else {"pii_redaction_config": config}), + ) + + context = {"thread_id": "memory-depth", "user_id": "memory-owner"} + result = graph.invoke({"messages": messages}, context=context) if mode == "sync" else await graph.ainvoke({"messages": messages}, context=context) + + assert result["messages"][-1].content == "Task complete." + assert json.dumps(attempt.tool_calls[0]["args"]) == original + selected, unused = (manager.add, manager.aadd) if mode == "sync" else (manager.aadd, manager.add) + unused.assert_not_called() + assert selected.call_count == int(memory) + if not memory: + return + call = selected.call_args + assert call.args[0] == "memory-depth" + assert call.kwargs["user_id"] == "memory-owner" + queued = next(message for message in call.args[1] if message.id == "historical-ai") + assert queued.tool_calls[0]["id"] == "old-call" + assert queued.tool_calls[0]["name"] == "_consume" + expected = {redact_text(_EMAIL, config): redact_text(_OTHER_EMAIL, config)} if policy == "enabled" else {_EMAIL: _OTHER_EMAIL} + assert _leaf(queued.tool_calls[0]["args"], depth) == {**expected, "number": 1, "flag": False, "empty": None} + assert (queued is attempt) is (policy != "enabled") + + +@pytest.mark.parametrize("shape,depth", [("dict", 1100), ("mixed", 550)]) +def test_deep_arguments_are_redacted_without_mutating_messages(shape, depth): + depth = _json_depth_limit(depth, shape) + args = _arguments(depth, shape) + original = json.dumps(args) + message = AIMessage(content="", id="ai-1", tool_calls=[{"name": "lookup", "args": args, "id": "call-1"}]) + config = _policy("enabled") + + queued = memory_middleware.redact_queued_messages([message], config)[0] + + assert queued.id == message.id + assert queued is not message + redacted = queued.tool_calls[0]["args"] + assert queued.tool_calls[0]["id"] == "call-1" + leaf = _leaf(redacted, depth, shape) + assert leaf == {redact_text(_EMAIL, config): redact_text(_OTHER_EMAIL, config), "number": 1, "flag": False, "empty": None} + assert json.dumps(args) == original + leaf["new"] = "queued only" + assert "new" not in _leaf(args, depth, shape) + + +@pytest.mark.parametrize("policy", ["omitted", "disabled"]) +def test_disabled_redaction_returns_an_independent_message_list(policy): + message = AIMessage(content="done", tool_calls=[{"name": "lookup", "args": _arguments(_json_depth_limit(1100)), "id": "call-1"}]) + messages = [message] + + queued = memory_middleware.redact_queued_messages(messages, _policy(policy)) + + assert queued is not messages + assert queued[0] is message + queued.clear() + assert messages == [message] + + +def test_redaction_keeps_shared_branches_and_key_collision_order(): + config = _policy("enabled") + token = redact_text(_EMAIL, config) + shared = {_EMAIL: _OTHER_EMAIL} + args = {"contacts": {_EMAIL: "first", token: "last"}, "branches": [shared, shared]} + message = AIMessage(content="", tool_calls=[{"name": "lookup", "args": args, "id": "call-1"}]) + + queued = memory_middleware.redact_queued_messages([message], config)[0].tool_calls[0]["args"] + + assert queued["contacts"] == {token: "last"} + assert queued["branches"] == [{token: redact_text(_OTHER_EMAIL, config)}] * 2 + queued["branches"][0]["new"] = "queued only" + assert shared == {_EMAIL: _OTHER_EMAIL} + assert "new" not in queued["branches"][1] + + +@pytest.mark.parametrize("shape", ["dict", "list"]) +def test_non_json_cycles_fail_without_an_unredacted_fallback(shape): + cyclic = {} if shape == "dict" else [] + if shape == "dict": + cyclic["cycle"] = cyclic + else: + cyclic.append(cyclic) + message = AIMessage(content="", tool_calls=[{"name": "lookup", "args": {"nested": cyclic, "email": _EMAIL}, "id": "call-1"}]) + + with pytest.raises(ValueError, match="[Cc]ycl"): + memory_middleware.redact_queued_messages([message], _policy("enabled")) + + +def test_deep_arguments_without_pii_matches_keep_message_identity(): + args = {"text": "ordinary text"} + for _ in range(_json_depth_limit(10_000)): + args = {"nested": args} + args = json.loads(json.dumps(args)) + message = AIMessage(content="done", tool_calls=[{"name": "lookup", "args": args, "id": "call-1"}]) + + queued = memory_middleware.redact_queued_messages([message], _policy("enabled")) + + assert queued[0] is message + + +@pytest.mark.parametrize("view", ["provider-raw", "legacy"]) +@pytest.mark.parametrize("redact_email", [False, True]) +def test_raw_argument_admission_at_json_depth_boundary(view, redact_email): + depth = _json_depth_limit(10_000) + args = _arguments(depth) + original = json.dumps(args) + raw_call = {"name": "lookup", "arguments": args} + # Restored messages can retain structured raw arguments separately from + # their parsed view. Both carriers must follow the same redaction policy. + field, raw = ("tool_calls", [{"id": "call-1", "type": "function", "function": raw_call}]) if view == "provider-raw" else ("function_call", raw_call) + message = AIMessage(content="done", id="raw-ai", tool_calls=[{"name": "lookup", "args": {}, "id": "call-1"}], additional_kwargs={field: raw}) + config = PiiRedactionConfig(enabled=True, redact_email=redact_email, token_secret="memory-depth-regression-secret") + + queued = memory_middleware.redact_queued_messages([message], config)[0] + + result = queued.additional_kwargs[field][0]["function"]["arguments"] if view == "provider-raw" else queued.additional_kwargs[field]["arguments"] + expected = {redact_text(_EMAIL, config): redact_text(_OTHER_EMAIL, config)} if redact_email else {_EMAIL: _OTHER_EMAIL} + assert _leaf(result, depth) == {**expected, "number": 1, "flag": False, "empty": None} + assert queued.id == message.id and queued.tool_calls[0]["id"] == "call-1" + assert (queued is message) is not redact_email + assert json.dumps(args) == original