mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-09 16:08:31 +00:00
fix(middleware): prevent ID-swap recursive injection and orphan peer compression (#3746)
Some checks are pending
Backend Blocking IO / backend-blocking-io (push) Waiting to run
Unit Tests / backend-unit-tests (push) Waiting to run
E2E Tests / e2e-tests (push) Waiting to run
Frontend Unit Tests / frontend-unit-tests (push) Waiting to run
Lint Check / lint-backend (push) Waiting to run
Lint Check / lint-frontend (push) Waiting to run
Replay E2E (front-back contract) / Layer 1 — backend golden (no API key) (push) Waiting to run
Replay E2E (front-back contract) / Layer 2 — full-stack render (no API key) (push) Waiting to run
Some checks are pending
Backend Blocking IO / backend-blocking-io (push) Waiting to run
Unit Tests / backend-unit-tests (push) Waiting to run
E2E Tests / e2e-tests (push) Waiting to run
Frontend Unit Tests / frontend-unit-tests (push) Waiting to run
Lint Check / lint-backend (push) Waiting to run
Lint Check / lint-frontend (push) Waiting to run
Replay E2E (front-back contract) / Layer 1 — backend golden (no API key) (push) Waiting to run
Replay E2E (front-back contract) / Layer 2 — full-stack render (no API key) (push) Waiting to run
- Add endswith('__user') guard in _is_user_injection_target to prevent
unbounded suffix growth (id__user__user__user...) when a prior ID-swap
peer is mistakenly treated as a new injection target
- Add peer rescue in _preserve_dynamic_context_reminders to keep
ID-swap __user and __memory messages out of summary compression,
preventing orphaned SystemMessage reminders and lost user context
- Add 6 regression tests covering both failure modes
This commit is contained in:
parent
820560e576
commit
f656440245
4 changed files with 374 additions and 5 deletions
|
|
@ -105,7 +105,21 @@ def _last_injected_date(messages: list) -> str | None:
|
|||
|
||||
def _is_user_injection_target(message: object) -> bool:
|
||||
"""Return whether *message* can receive a dynamic-context reminder."""
|
||||
return isinstance(message, HumanMessage) and not is_dynamic_context_reminder(message) and message.name != _SUMMARY_MESSAGE_NAME
|
||||
if not isinstance(message, HumanMessage):
|
||||
return False
|
||||
if is_dynamic_context_reminder(message):
|
||||
return False
|
||||
if message.name == _SUMMARY_MESSAGE_NAME:
|
||||
return False
|
||||
# Prevent recursive ID-swap: a message whose ID ends with "__user" was
|
||||
# produced by a prior _make_reminder_and_user_messages call and must not
|
||||
# be processed again — doing so causes unbounded suffix growth
|
||||
# (id__user__user__user...) and ghost-message re-execution.
|
||||
# Using endswith (not substring "in") avoids false positives on IDs that
|
||||
# happen to contain "__user" in the middle.
|
||||
if message.id and str(message.id).endswith("__user"):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class DynamicContextMiddleware(AgentMiddleware):
|
||||
|
|
|
|||
|
|
@ -256,18 +256,54 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
|
|||
messages_to_summarize: list[AnyMessage],
|
||||
preserved_messages: list[AnyMessage],
|
||||
) -> tuple[list[AnyMessage], list[AnyMessage]]:
|
||||
"""Keep hidden dynamic-context reminders out of summary compression.
|
||||
"""Keep hidden dynamic-context reminders and their ID-swap peers out of summary compression.
|
||||
|
||||
These reminders carry the current date and optional memory. If summarization
|
||||
removes them, DynamicContextMiddleware can mistake the summary HumanMessage
|
||||
for the first user message and inject the reminder in the wrong place.
|
||||
|
||||
The ID-swap triplet produced by ``_make_reminder_and_user_messages`` contains
|
||||
three messages: ``SystemMessage(id=X)`` and ``HumanMessage(id=X__memory)`` are
|
||||
both tagged with ``dynamic_context_reminder=True``, but ``HumanMessage(id=X__user)``
|
||||
carries the original user content and is **not** tagged. Without peer rescue,
|
||||
``__user`` would stay in ``to_summarize`` and be compressed into prose — orphaning
|
||||
the tagged messages and losing the user question from the model's direct context.
|
||||
|
||||
This method rescues tagged reminders and also rescues any untagged messages whose
|
||||
``id`` shares the same ``stable_id`` prefix (i.e. ``X__user``, ``X__memory``).
|
||||
"""
|
||||
reminders = [msg for msg in messages_to_summarize if is_dynamic_context_reminder(msg)]
|
||||
if not reminders:
|
||||
return messages_to_summarize, preserved_messages
|
||||
|
||||
remaining = [msg for msg in messages_to_summarize if not is_dynamic_context_reminder(msg)]
|
||||
return remaining, reminders + preserved_messages
|
||||
# Collect the base IDs (the stable_id prefix) from tagged reminders.
|
||||
# For a reminder with id="ctx-001__memory", the base is "ctx-001".
|
||||
# For a reminder with id="ctx-001" (SystemMessage), the base is "ctx-001".
|
||||
# removesuffix is suffix-only — it won't strip a "__" that sits in the
|
||||
# middle of a stable_id (e.g. "ctx__001" stays intact, unlike rsplit
|
||||
# which would mis-derive "ctx"). Only known ID-swap suffixes (__memory,
|
||||
# __user) are stripped; __user is not tagged so won't appear in reminders,
|
||||
# but is included defensively.
|
||||
reminder_base_ids: set[str] = set()
|
||||
for msg in reminders:
|
||||
if msg.id:
|
||||
base = msg.id.removesuffix("__memory").removesuffix("__user")
|
||||
reminder_base_ids.add(base)
|
||||
|
||||
# Single-pass partition: walk messages_to_summarize in chronological order
|
||||
# and rescue both tagged reminders and untagged ID-swap peers (whose id
|
||||
# starts with a known base + "__"). This preserves the original message
|
||||
# order within rescued — critical when multiple triplets land in one
|
||||
# summarization window — and eliminates the need for id(m)-based dedup
|
||||
# that the previous reminders+peers concatenation required.
|
||||
rescued: list[AnyMessage] = []
|
||||
remaining: list[AnyMessage] = []
|
||||
for msg in messages_to_summarize:
|
||||
if is_dynamic_context_reminder(msg) or (msg.id and any(msg.id.startswith(b + "__") for b in reminder_base_ids)):
|
||||
rescued.append(msg)
|
||||
else:
|
||||
remaining.append(msg)
|
||||
return remaining, rescued + preserved_messages
|
||||
|
||||
def _partition_with_skill_rescue(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -518,3 +518,117 @@ def test_no_second_midnight_injection_once_date_updated():
|
|||
result = mw.before_agent(state, _fake_runtime())
|
||||
|
||||
assert result is None # same day as last injected date → no update
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ID-swap recursive-injection guard (issue #3725)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_user_suffix_message_is_not_injection_target():
|
||||
"""Regression guard: HumanMessage whose ID ends with ``__user`` must not be
|
||||
treated as an injection target.
|
||||
|
||||
After the ID-swap in ``_make_reminder_and_user_messages``, the original user
|
||||
text becomes ``HumanMessage(id=X__user)``. If the middleware processes this
|
||||
message again, it would perform another ID-swap → ``X__user__user`` → … →
|
||||
unbounded suffix growth and ghost-message re-execution (issue #3725).
|
||||
"""
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import _is_user_injection_target
|
||||
|
||||
# A __user-suffix message is NOT a valid injection target
|
||||
user_swap_msg = HumanMessage(content="Hello", id="msg-1__user")
|
||||
assert _is_user_injection_target(user_swap_msg) is False
|
||||
|
||||
# A __memory-suffix message is already tagged as a reminder, also rejected
|
||||
memory_swap_msg = HumanMessage(
|
||||
content="<memory>prefs</memory>",
|
||||
id="msg-1__memory",
|
||||
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
)
|
||||
assert _is_user_injection_target(memory_swap_msg) is False
|
||||
|
||||
# A normal HumanMessage without __user suffix IS a valid target
|
||||
normal_msg = HumanMessage(content="Hello", id="msg-1")
|
||||
assert _is_user_injection_target(normal_msg) is True
|
||||
|
||||
|
||||
def test_endswith_not_substring_prevents_false_positive():
|
||||
"""``endswith("__user")`` must NOT reject messages whose ID merely contains
|
||||
``__user`` somewhere in the middle (e.g. ``user__question-123``).
|
||||
|
||||
A substring check (``"__user" in id``) would incorrectly reject such IDs.
|
||||
"""
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import _is_user_injection_target
|
||||
|
||||
# ID contains "__user" in the middle — should NOT be rejected
|
||||
middle_match = HumanMessage(content="question", id="user__question-123")
|
||||
assert _is_user_injection_target(middle_match) is True
|
||||
|
||||
# ID ends with "__user" — should be rejected
|
||||
suffix_match = HumanMessage(content="question", id="msg-1__user")
|
||||
assert _is_user_injection_target(suffix_match) is False
|
||||
|
||||
# Nested suffix "__user__user" — should also be rejected (recursive case)
|
||||
recursive_match = HumanMessage(content="question", id="msg-1__user__user")
|
||||
assert _is_user_injection_target(recursive_match) is False
|
||||
|
||||
|
||||
def test_no_recursive_id_swap_in_full_middleware_flow():
|
||||
"""End-to-end guard: after the first ID-swap, a second call to ``before_agent``
|
||||
must NOT produce a second swap on the ``__user`` message.
|
||||
|
||||
This reproduces the exact scenario from issue #3725: a session with an
|
||||
existing ID-swap triplet receives a new HumanMessage, and the middleware
|
||||
must only inject into the new message — not re-process the ``__user`` peer.
|
||||
|
||||
The state_v2 reminder deliberately omits the parseable date from both
|
||||
content and additional_kwargs so ``_last_injected_date`` returns None.
|
||||
This forces the first-turn injection path to actually reach
|
||||
``_is_user_injection_target``, which must reject ``msg-1__user`` and
|
||||
select ``msg-2`` instead — exercising the endswith("__user") guard
|
||||
end-to-end rather than relying on the same-day short-circuit.
|
||||
"""
|
||||
mw = _make_middleware()
|
||||
|
||||
# First call: inject into HumanMessage(id="msg-1")
|
||||
state_v1 = {"messages": [HumanMessage(content="Hello", id="msg-1")]}
|
||||
|
||||
with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt, mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""):
|
||||
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
|
||||
result_v1 = mw.before_agent(state_v1, _fake_runtime())
|
||||
|
||||
assert result_v1 is not None
|
||||
msgs_v1 = result_v1["messages"]
|
||||
assert len(msgs_v1) == 2
|
||||
assert msgs_v1[0].id == "msg-1" # reminder takes original ID
|
||||
assert msgs_v1[1].id == "msg-1__user" # user content gets derived ID
|
||||
|
||||
# Simulate state after first turn: ID-swap triplet (without parseable date
|
||||
# so _last_injected_date returns None → first-turn path is exercised)
|
||||
# + AI reply + new user message.
|
||||
state_v2 = {
|
||||
"messages": [
|
||||
SystemMessage(
|
||||
content="<system-reminder>\nplaceholder\n</system-reminder>",
|
||||
id="msg-1",
|
||||
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
),
|
||||
HumanMessage(content="Hello", id="msg-1__user"),
|
||||
AIMessage(content="Hi there"),
|
||||
HumanMessage(content="Follow-up", id="msg-2"),
|
||||
]
|
||||
}
|
||||
|
||||
# Second call: _last_injected_date returns None (no parseable date),
|
||||
# so _inject enters first-turn path and must skip msg-1__user via the
|
||||
# endswith("__user") guard, then inject into msg-2.
|
||||
with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt, mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""):
|
||||
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
|
||||
result_v2 = mw.before_agent(state_v2, _fake_runtime())
|
||||
|
||||
# The guard must route injection to msg-2, not msg-1__user.
|
||||
assert result_v2 is not None
|
||||
msgs_v2 = result_v2["messages"]
|
||||
assert msgs_v2[0].id == "msg-2" # reminder takes new message's ID
|
||||
assert msgs_v2[1].id == "msg-2__user" # user content gets derived ID
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ from langchain_core.outputs import ChatGeneration, ChatResult
|
|||
from langgraph.constants import TAG_NOSTREAM
|
||||
|
||||
from deerflow.agents.memory.summarization_hook import memory_flush_hook
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, DynamicContextMiddleware
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, DynamicContextMiddleware, is_dynamic_context_reminder
|
||||
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware, SummarizationEvent
|
||||
from deerflow.config.memory_config import MemoryConfig
|
||||
|
||||
|
|
@ -822,3 +822,208 @@ def test_memory_flush_hook_passes_runtime_user_id(monkeypatch: pytest.MonkeyPatc
|
|||
|
||||
queue.add_nowait.assert_called_once()
|
||||
assert queue.add_nowait.call_args.kwargs["user_id"] == "alice"
|
||||
|
||||
|
||||
def test_id_swap_user_peer_is_preserved_across_summarization() -> None:
|
||||
"""__user (untagged) must be rescued alongside its tagged ID-swap peers.
|
||||
|
||||
The ID-swap triplet from _make_reminder_and_user_messages is:
|
||||
[SystemMessage(id=X, reminder=True), HumanMessage(id=X__memory, reminder=True),
|
||||
HumanMessage(id=X__user)] — only the first two are tagged. Without peer
|
||||
rescue, __user stays in to_summarize and is compressed into prose, orphaning
|
||||
the tagged messages and losing the user question from direct model context.
|
||||
"""
|
||||
captured: list[SummarizationEvent] = []
|
||||
middleware = _middleware(before_summarization=[captured.append])
|
||||
|
||||
# Build an ID-swap triplet (SystemMessage + __memory + __user)
|
||||
stable_id = "ctx-001"
|
||||
reminder_system = SystemMessage(
|
||||
content="<system-reminder>\n<current_date>2026-05-08, Friday</current_date>\n</system-reminder>",
|
||||
id=stable_id,
|
||||
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
)
|
||||
memory_msg = HumanMessage(
|
||||
content="<memory>user preferences</memory>",
|
||||
id=f"{stable_id}__memory",
|
||||
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
)
|
||||
user_msg = HumanMessage(
|
||||
content="What is the weather in Tokyo?",
|
||||
id=f"{stable_id}__user",
|
||||
)
|
||||
|
||||
result = middleware.before_model(
|
||||
{
|
||||
"messages": [
|
||||
reminder_system,
|
||||
memory_msg,
|
||||
user_msg,
|
||||
AIMessage(content="The weather is sunny.", id="ai-1"),
|
||||
HumanMessage(content="user-2"),
|
||||
]
|
||||
},
|
||||
_runtime(),
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
# The __user message should NOT be in messages_to_summarize
|
||||
summarized_contents = [m.content for m in captured[0].messages_to_summarize]
|
||||
assert "What is the weather in Tokyo?" not in summarized_contents
|
||||
|
||||
# All three triplet members should be in preserved_messages
|
||||
preserved_ids = [m.id for m in captured[0].preserved_messages]
|
||||
assert stable_id in preserved_ids
|
||||
assert f"{stable_id}__memory" in preserved_ids
|
||||
assert f"{stable_id}__user" in preserved_ids
|
||||
|
||||
# The emitted state includes all three triplet members
|
||||
emitted = result["messages"]
|
||||
assert isinstance(emitted[0], RemoveMessage)
|
||||
# Find the triplet members in the emitted messages
|
||||
emitted_ids = [m.id for m in emitted[2:]] # Skip RemoveMessage + summary
|
||||
assert stable_id in emitted_ids
|
||||
assert f"{stable_id}__memory" in emitted_ids
|
||||
assert f"{stable_id}__user" in emitted_ids
|
||||
|
||||
|
||||
def test_id_swap_user_peer_preserved_without_memory() -> None:
|
||||
"""When there's no __memory in the triplet, __user is still rescued."""
|
||||
captured: list[SummarizationEvent] = []
|
||||
middleware = _middleware(before_summarization=[captured.append])
|
||||
|
||||
stable_id = "ctx-002"
|
||||
reminder_system = SystemMessage(
|
||||
content="<system-reminder>\n<current_date>2026-05-09, Saturday</current_date>\n</system-reminder>",
|
||||
id=stable_id,
|
||||
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
)
|
||||
user_msg = HumanMessage(
|
||||
content="How are you?",
|
||||
id=f"{stable_id}__user",
|
||||
)
|
||||
|
||||
middleware.before_model(
|
||||
{
|
||||
"messages": [
|
||||
reminder_system,
|
||||
user_msg,
|
||||
AIMessage(content="I'm fine.", id="ai-2"),
|
||||
HumanMessage(content="user-3"),
|
||||
]
|
||||
},
|
||||
_runtime(),
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
summarized_contents = [m.content for m in captured[0].messages_to_summarize]
|
||||
assert "How are you?" not in summarized_contents
|
||||
|
||||
preserved_ids = [m.id for m in captured[0].preserved_messages]
|
||||
assert stable_id in preserved_ids
|
||||
assert f"{stable_id}__user" in preserved_ids
|
||||
|
||||
|
||||
def test_non_reminder_messages_with_double_underscore_id_not_rescued() -> None:
|
||||
"""Messages whose IDs contain "__" but are NOT ID-swap peers are not rescued."""
|
||||
captured: list[SummarizationEvent] = []
|
||||
middleware = _middleware(before_summarization=[captured.append])
|
||||
|
||||
# A normal reminder without any ID-swap peers
|
||||
reminder = _dynamic_context_reminder("standalone-reminder")
|
||||
# A message whose ID happens to contain "__" but is unrelated
|
||||
unrelated = HumanMessage(content="unrelated question", id="some-other__msg")
|
||||
|
||||
middleware.before_model(
|
||||
{
|
||||
"messages": [
|
||||
reminder,
|
||||
unrelated,
|
||||
AIMessage(content="answer"),
|
||||
HumanMessage(content="user-2"),
|
||||
]
|
||||
},
|
||||
_runtime(),
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
# The unrelated message is NOT rescued — it stays in to_summarize
|
||||
preserved_ids = [m.id for m in captured[0].preserved_messages]
|
||||
assert "some-other__msg" not in preserved_ids
|
||||
# Only the standalone reminder is rescued (no peer lookup triggered)
|
||||
assert "standalone-reminder" in preserved_ids
|
||||
|
||||
|
||||
def test_multiple_id_swap_triplets_preserve_chronological_order() -> None:
|
||||
"""When multiple ID-swap triplets sit in one summarization window, rescued
|
||||
messages must retain their original chronological order — not be scrambled
|
||||
by separating tagged reminders from untagged peers.
|
||||
|
||||
Regression: the previous reminders+peers concatenation rescued as
|
||||
[Sys(base1), Sys(base2), Mem(base1), Mem(base2), User(base1), User(base2)],
|
||||
detaching each user question from its AI answer. The single-pass partition
|
||||
preserves [Sys(base1), Mem(base1), User(base1), Sys(base2), Mem(base2), User(base2)].
|
||||
"""
|
||||
captured: list[SummarizationEvent] = []
|
||||
middleware = _middleware(before_summarization=[captured.append])
|
||||
|
||||
# Two complete triplets (first-turn + midnight crossing) plus an AI reply
|
||||
# between them, all sitting before the summarization cutoff.
|
||||
base1 = "ctx-001"
|
||||
base2 = "ctx-002"
|
||||
reminder_1 = SystemMessage(
|
||||
content="<system-reminder>\n<current_date>2026-05-08, Friday</current_date>\n</system-reminder>",
|
||||
id=base1,
|
||||
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
)
|
||||
memory_1 = HumanMessage(
|
||||
content="<memory>prefs v1</memory>",
|
||||
id=f"{base1}__memory",
|
||||
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
)
|
||||
user_1 = HumanMessage(content="What is the weather?", id=f"{base1}__user")
|
||||
ai_1 = AIMessage(content="Sunny.", id="ai-1")
|
||||
|
||||
reminder_2 = SystemMessage(
|
||||
content="<system-reminder>\n<current_date>2026-05-09, Saturday</current_date>\n</system-reminder>",
|
||||
id=base2,
|
||||
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
)
|
||||
memory_2 = HumanMessage(
|
||||
content="<memory>prefs v2</memory>",
|
||||
id=f"{base2}__memory",
|
||||
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
|
||||
)
|
||||
user_2 = HumanMessage(content="How are you?", id=f"{base2}__user")
|
||||
ai_2 = AIMessage(content="Fine.", id="ai-2")
|
||||
|
||||
middleware.before_model(
|
||||
{
|
||||
"messages": [
|
||||
reminder_1,
|
||||
memory_1,
|
||||
user_1,
|
||||
ai_1,
|
||||
reminder_2,
|
||||
memory_2,
|
||||
user_2,
|
||||
ai_2,
|
||||
HumanMessage(content="latest question"),
|
||||
]
|
||||
},
|
||||
_runtime(),
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
# Rescued messages must appear in their original chronological order:
|
||||
# each triplet stays contiguous, not re-grouped by role.
|
||||
preserved = captured[0].preserved_messages
|
||||
rescued_ids = [m.id for m in preserved if m.id and (is_dynamic_context_reminder(m) or m.id in (f"{base1}__user", f"{base2}__user"))]
|
||||
assert rescued_ids == [
|
||||
base1,
|
||||
f"{base1}__memory",
|
||||
f"{base1}__user",
|
||||
base2,
|
||||
f"{base2}__memory",
|
||||
f"{base2}__user",
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue