mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-08 03:30:34 +00:00
7e64330dd1
* fix(memory): avoid recursive tool-argument admission * fix(memory): track argument redaction changes without deep equality
213 lines
9.4 KiB
Python
213 lines
9.4 KiB
Python
"""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
|