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

- 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:
Zhipeng Zheng 2026-06-24 17:55:26 +08:00 committed by GitHub
parent 820560e576
commit f656440245
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 374 additions and 5 deletions

View file

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

View file

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

View file

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

View file

@ -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",
]