From f65644024586d6ecb7b20e33ba7611f9c9e49dc3 Mon Sep 17 00:00:00 2001 From: Zhipeng Zheng <2791184972@qq.com> Date: Wed, 24 Jun 2026 17:55:26 +0800 Subject: [PATCH] fix(middleware): prevent ID-swap recursive injection and orphan peer compression (#3746) - 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 --- .../middlewares/dynamic_context_middleware.py | 16 +- .../middlewares/summarization_middleware.py | 42 +++- .../tests/test_dynamic_context_middleware.py | 114 ++++++++++ .../tests/test_summarization_middleware.py | 207 +++++++++++++++++- 4 files changed, 374 insertions(+), 5 deletions(-) diff --git a/backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py index 6b3ebb6ad..726a6a61a 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py @@ -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): diff --git a/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py index 917823612..4c89ee72d 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py @@ -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, diff --git a/backend/tests/test_dynamic_context_middleware.py b/backend/tests/test_dynamic_context_middleware.py index 71bc5cf24..49dcb4931 100644 --- a/backend/tests/test_dynamic_context_middleware.py +++ b/backend/tests/test_dynamic_context_middleware.py @@ -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="prefs", + 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="\nplaceholder\n", + 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 diff --git a/backend/tests/test_summarization_middleware.py b/backend/tests/test_summarization_middleware.py index 0ab526d4f..8b70429f8 100644 --- a/backend/tests/test_summarization_middleware.py +++ b/backend/tests/test_summarization_middleware.py @@ -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="\n2026-05-08, Friday\n", + id=stable_id, + additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True}, + ) + memory_msg = HumanMessage( + content="user preferences", + 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="\n2026-05-09, Saturday\n", + 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="\n2026-05-08, Friday\n", + id=base1, + additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True}, + ) + memory_1 = HumanMessage( + content="prefs v1", + 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="\n2026-05-09, Saturday\n", + id=base2, + additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True}, + ) + memory_2 = HumanMessage( + content="prefs v2", + 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", + ]