fix(memory): avoid recursive tool-argument admission (#7928)

* fix(memory): avoid recursive tool-argument admission

* fix(memory): track argument redaction changes without deep equality
This commit is contained in:
lilei 2026-11-06 21:24:57 +08:00 committed by GitHub
parent a41d0fcec1
commit 7e64330dd1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 274 additions and 21 deletions

View File

@ -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.

View File

@ -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`.

View File

@ -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):

View File

@ -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