fix(harness): preserve dynamic context across summarization (#2823)

This commit is contained in:
DanielWalnut 2026-05-09 19:39:36 +08:00 committed by GitHub
parent f76e4e35c8
commit 881ff71252
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 100 additions and 2 deletions

View File

@ -45,6 +45,7 @@ logger = logging.getLogger(__name__)
_DATE_RE = re.compile(r"<current_date>([^<]+)</current_date>")
_DYNAMIC_CONTEXT_REMINDER_KEY = "dynamic_context_reminder"
_SUMMARY_MESSAGE_NAME = "summary"
def _extract_date(content: str) -> str | None:
@ -72,6 +73,16 @@ def _last_injected_date(messages: list) -> str | None:
return None
def is_dynamic_context_reminder(message: object) -> bool:
"""Return whether *message* is a hidden dynamic-context reminder."""
return isinstance(message, HumanMessage) and bool(message.additional_kwargs.get(_DYNAMIC_CONTEXT_REMINDER_KEY))
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
class DynamicContextMiddleware(AgentMiddleware):
"""Inject memory and current date into HumanMessages as a <system-reminder>.
@ -163,7 +174,7 @@ class DynamicContextMiddleware(AgentMiddleware):
if last_date is None:
# ── First turn: inject full reminder as a separate HumanMessage ─────
first_idx = next((i for i, m in enumerate(messages) if isinstance(m, HumanMessage)), None)
first_idx = next((i for i, m in enumerate(messages) if _is_user_injection_target(m)), None)
if first_idx is None:
return None
full_reminder = self._build_full_reminder()
@ -181,7 +192,7 @@ class DynamicContextMiddleware(AgentMiddleware):
return None
# ── Midnight crossed: inject date-update reminder as a separate HumanMessage ──
last_human_idx = next((i for i in reversed(range(len(messages))) if isinstance(messages[i], HumanMessage)), None)
last_human_idx = next((i for i in reversed(range(len(messages))) if _is_user_injection_target(messages[i])), None)
if last_human_idx is None:
return None

View File

@ -14,6 +14,7 @@ from langgraph.config import get_config
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langgraph.runtime import Runtime
from deerflow.agents.middlewares.dynamic_context_middleware import is_dynamic_context_reminder
from deerflow.agents.middlewares.tool_call_metadata import clone_ai_message_with_tool_calls
logger = logging.getLogger(__name__)
@ -135,6 +136,7 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
return None
messages_to_summarize, preserved_messages = self._partition_with_skill_rescue(messages, cutoff_index)
messages_to_summarize, preserved_messages = self._preserve_dynamic_context_reminders(messages_to_summarize, preserved_messages)
self._fire_hooks(messages_to_summarize, preserved_messages, runtime)
summary = self._create_summary(messages_to_summarize)
new_messages = self._build_new_messages(summary)
@ -160,6 +162,7 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
return None
messages_to_summarize, preserved_messages = self._partition_with_skill_rescue(messages, cutoff_index)
messages_to_summarize, preserved_messages = self._preserve_dynamic_context_reminders(messages_to_summarize, preserved_messages)
self._fire_hooks(messages_to_summarize, preserved_messages, runtime)
summary = await self._acreate_summary(messages_to_summarize)
new_messages = self._build_new_messages(summary)
@ -179,6 +182,24 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
"""
return [HumanMessage(content=f"Here is a summary of the conversation to date:\n\n{summary}", name="summary")]
def _preserve_dynamic_context_reminders(
self,
messages_to_summarize: list[AnyMessage],
preserved_messages: list[AnyMessage],
) -> tuple[list[AnyMessage], list[AnyMessage]]:
"""Keep hidden dynamic-context reminders 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.
"""
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
def _partition_with_skill_rescue(
self,
messages: list[AnyMessage],

View File

@ -139,6 +139,30 @@ def test_injects_only_into_first_human_message_not_later_ones():
assert all(m.id != "msg-2" for m in msgs)
def test_summary_human_message_is_not_used_as_injection_target():
"""After summarization, the synthetic summary HumanMessage is not a user turn."""
mw = _make_middleware()
state = {
"messages": [
HumanMessage(content="Here is a summary of the conversation to date:\n\n...", id="summary-1", name="summary"),
AIMessage(content="Earlier reply"),
HumanMessage(content="Follow-up", id="msg-2"),
]
}
with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value=""), mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
result = mw.before_agent(state, _fake_runtime())
assert result is not None
msgs = result["messages"]
assert len(msgs) == 2
assert msgs[0].id == "msg-2"
assert msgs[0].additional_kwargs.get(_DYNAMIC_CONTEXT_REMINDER_KEY) is True
assert msgs[1].id == "msg-2__user"
assert msgs[1].content == "Follow-up"
# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------

View File

@ -1,12 +1,14 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest import mock
from unittest.mock import MagicMock
import pytest
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, ToolMessage
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.summarization_middleware import DeerFlowSummarizationMiddleware, SummarizationEvent
from deerflow.config.memory_config import MemoryConfig
@ -20,6 +22,14 @@ def _messages() -> list:
]
def _dynamic_context_reminder(msg_id: str = "reminder-1") -> HumanMessage:
return HumanMessage(
content="<system-reminder>\n<current_date>2026-05-08, Friday</current_date>\n</system-reminder>",
id=msg_id,
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
)
def _runtime(thread_id: str | None = "thread-1", agent_name: str | None = None) -> SimpleNamespace:
context = {}
if thread_id is not None:
@ -98,6 +108,38 @@ def test_before_summarization_hook_receives_messages_before_compression() -> Non
assert result["messages"][1].content.startswith("Here is a summary")
def test_dynamic_context_reminder_is_preserved_across_summarization() -> None:
captured: list[SummarizationEvent] = []
middleware = _middleware(before_summarization=[captured.append])
reminder = _dynamic_context_reminder()
result = middleware.before_model(
{
"messages": [
reminder,
HumanMessage(content="user-1"),
AIMessage(content="assistant-1"),
HumanMessage(content="user-2"),
]
},
_runtime(),
)
assert len(captured) == 1
assert [message.content for message in captured[0].messages_to_summarize] == ["user-1"]
assert captured[0].preserved_messages[0] is reminder
emitted = result["messages"]
assert isinstance(emitted[0], RemoveMessage)
assert emitted[1].name == "summary"
assert emitted[2] is reminder
followup_state = {"messages": [*emitted[1:], HumanMessage(content="Follow-up", id="msg-2")]}
with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt:
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
assert DynamicContextMiddleware().before_agent(followup_state, _runtime()) is None
def test_before_summarization_hook_not_called_when_threshold_not_met() -> None:
captured: list[SummarizationEvent] = []
middleware = _middleware(before_summarization=[captured.append], trigger=("messages", 10))