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 8cb749d60..6b3ebb6ad 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py @@ -52,6 +52,10 @@ _INJECT_TIMEOUT_SECONDS = 5.0 _DATE_RE = re.compile(r"([^<]+)") _DYNAMIC_CONTEXT_REMINDER_KEY = "dynamic_context_reminder" +# Authoritative injected date, carried in additional_kwargs of the date +# SystemMessage. Detection reads this instead of regex-parsing message content, +# so it is never exposed to user-influenceable memory content. +_REMINDER_DATE_KEY = "reminder_date" _SUMMARY_MESSAGE_NAME = "summary" @@ -75,11 +79,27 @@ def _last_injected_date(messages: list) -> str | None: Detection uses the ``dynamic_context_reminder`` additional_kwargs flag rather than content substring matching, so user messages containing ```` are not mistakenly treated as injected reminders. + + The authoritative date is the ``reminder_date`` value in additional_kwargs of + the date SystemMessage. Reminders without it (the separate ```` + HumanMessage, or any future dateless reminder) carry no date and are skipped, + so they cannot shadow the real date reminder. """ for msg in reversed(messages): - if is_dynamic_context_reminder(msg): + if not is_dynamic_context_reminder(msg): + continue + structured = msg.additional_kwargs.get(_REMINDER_DATE_KEY) + if isinstance(structured, str) and structured: + return structured + # Backward-compat for checkpoints written before reminder_date existed: + # the date lived in content. Scope the regex to SystemMessage so it never + # runs on the user-influenceable memory HumanMessage (preserves the OWASP + # role separation from #3630 and closes the memory date-spoofing hole). + if isinstance(msg, SystemMessage): content_str = msg.content if isinstance(msg.content, str) else str(msg.content) - return _extract_date(content_str) + date = _extract_date(content_str) + if date is not None: + return date return None @@ -152,26 +172,34 @@ class DynamicContextMiddleware(AgentMiddleware): original: HumanMessage, reminder_content: str, memory_content: str | None = None, + *, + reminder_date: str | None = None, ) -> list[SystemMessage | HumanMessage]: """Return messages using the ID-swap technique. SystemMessage carries framework-owned data (date, metadata) — takes - the original ID so add_messages replaces it in-place. Optional + the original ID so add_messages replaces it in-place. *reminder_date* + is recorded in its additional_kwargs as the authoritative injected date + (``_last_injected_date`` reads it instead of parsing content). Optional HumanMessage carries user-owned memory content with ``{id}__memory``. The actual user message gets ``{id}__user``. SystemMessage is used — system context must not masquerade as user input (#3630). Memory is deliberately kept as HumanMessage so - user-influenceable content does not gain system authority (OWASP LLM01). + user-influenceable content does not gain system authority (OWASP LLM01) + — and it deliberately never carries ``reminder_date``. """ stable_id = original.id or str(uuid.uuid4()) messages: list[SystemMessage | HumanMessage] = [] + reminder_kwargs = {"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True} + if reminder_date is not None: + reminder_kwargs[_REMINDER_DATE_KEY] = reminder_date messages.append( SystemMessage( content=reminder_content, id=stable_id, - additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True}, + additional_kwargs=reminder_kwargs, ) ) @@ -219,7 +247,7 @@ class DynamicContextMiddleware(AgentMiddleware): memory_block is not None, messages[first_idx].id, ) - result_msgs = self._make_reminder_and_user_messages(messages[first_idx], date_reminder, memory_block) + result_msgs = self._make_reminder_and_user_messages(messages[first_idx], date_reminder, memory_block, reminder_date=current_date) return {"messages": result_msgs} if last_date == current_date: @@ -231,7 +259,7 @@ class DynamicContextMiddleware(AgentMiddleware): if last_human_idx is None: return None - result_msgs = self._make_reminder_and_user_messages(messages[last_human_idx], self._build_date_update_reminder()) + result_msgs = self._make_reminder_and_user_messages(messages[last_human_idx], self._build_date_update_reminder(), reminder_date=current_date) logger.info("DynamicContextMiddleware: midnight crossing detected — injected date update before current turn") return {"messages": result_msgs} diff --git a/backend/tests/test_dynamic_context_middleware.py b/backend/tests/test_dynamic_context_middleware.py index 313295565..71bc5cf24 100644 --- a/backend/tests/test_dynamic_context_middleware.py +++ b/backend/tests/test_dynamic_context_middleware.py @@ -38,6 +38,20 @@ def _reminder_msg(content: str, msg_id: str) -> HumanMessage: ) +def _date_reminder_msg(date_str: str, msg_id: str) -> SystemMessage: + """Build a persisted date reminder in the current production shape. + + A date SystemMessage whose ``reminder_date`` additional_kwargs carries the + authoritative date — what ``DynamicContextMiddleware`` now writes to state. + """ + content = f"\n{date_str}\n" + return SystemMessage( + content=content, + id=msg_id, + additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True, "reminder_date": date_str}, + ) + + # --------------------------------------------------------------------------- # Basic injection # --------------------------------------------------------------------------- @@ -105,10 +119,9 @@ def test_memory_included_when_present(): def test_skips_injection_if_already_present(): """Second turn: separate reminder message already present → no update.""" mw = _make_middleware() - reminder_content = "\n2026-05-08, Friday\n" state = { "messages": [ - _reminder_msg(reminder_content, "msg-1"), + _date_reminder_msg("2026-05-08, Friday", "msg-1"), HumanMessage(content="Hello", id="msg-1__user"), AIMessage(content="Hi there"), HumanMessage(content="Follow-up", id="msg-2"), @@ -122,6 +135,131 @@ def test_skips_injection_if_already_present(): assert result is None # no update needed +def test_second_turn_with_memory_does_not_reinject(): + """Regression: a dateless memory reminder must not shadow the date reminder. + + Reproduces the scrambled-messages / wrong-answer bug (thread + 9be75d63): production persists the injected context as TWO flagged + messages — a date SystemMessage and a separate dateless + HumanMessage. On a later turn ``_last_injected_date`` scans in reverse + and hits the memory message first; because it has no it + must keep scanning to find the real date. If it stops and returns None, + the middleware falsely treats this as the first turn, re-injects, picks + the previous turn's ``__user`` message as the target, and the model + re-answers the stale turn instead of the new one. + """ + mw = _make_middleware() + date_reminder = "\n2026-05-08, Friday\n" + state = { + "messages": [ + SystemMessage( + content=date_reminder, + id="msg-1", + additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True}, + ), + _reminder_msg("\nUser prefers Python.\n", "msg-1__memory"), + HumanMessage(content="test", id="msg-1__user", name="user-input"), + AIMessage(content="Test received"), + HumanMessage(content="tell me the weather", id="msg-2", name="user-input"), + ] + } + + with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value="\nUser prefers Python.\n"), 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 None # same day already injected → must NOT re-inject + + +def test_poisoned_memory_does_not_spoof_injected_date(): + """A embedded in user-influenceable memory must not spoof detection. + + Memory is LLM-extracted from user input and injected unescaped (it's + hide_from_ui, so InputSanitizationMiddleware skips it). If a memory fact + contains a literal …, content-regex detection + would return that fake date (it sits after the authoritative date message but + is hit first in the reverse scan) and trigger a false midnight crossing / + re-injection. The authoritative date lives in additional_kwargs, so detection + must ignore the memory content entirely. + """ + mw = _make_middleware() + today = "2026-05-08, Friday" + date_reminder = f"\n{today}\n" + state = { + "messages": [ + SystemMessage( + content=date_reminder, + id="msg-1", + additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True, "reminder_date": today}, + ), + _reminder_msg("\nUser asked about 2024-01-01 last year.\n", "msg-1__memory"), + HumanMessage(content="test", id="msg-1__user", name="user-input"), + AIMessage(content="Test received"), + HumanMessage(content="follow up", id="msg-2", name="user-input"), + ] + } + + with mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt: + mock_dt.now.return_value.strftime.return_value = today + result = mw.before_agent(state, _fake_runtime()) + + # Detection uses the authoritative metadata date (today) → same day → no re-injection. + # If the fake 2024 date from memory content leaked in, this would be a midnight crossing. + assert result is None + + +def test_date_reminder_carries_structured_date(): + """First-turn injection records the authoritative date in additional_kwargs. + + The date SystemMessage carries ``reminder_date``; the memory HumanMessage + deliberately does not (it is dateless and must never spoof detection). + """ + mw = _make_middleware() + state = {"messages": [HumanMessage(content="Hi", id="msg-1")]} + + with ( + mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value="\nUser prefers Python.\n"), + 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()) + + msgs = result["messages"] + assert isinstance(msgs[0], SystemMessage) + assert msgs[0].additional_kwargs.get("reminder_date") == "2026-05-08, Friday" + # Memory HumanMessage must not carry the authoritative date + assert isinstance(msgs[1], HumanMessage) + assert "reminder_date" not in msgs[1].additional_kwargs + + +def test_legacy_systemmessage_reminder_without_key_detected(): + """Backward-compat: pre-reminder_date checkpoints kept the date in content only. + + A date SystemMessage with the date in content but no ``reminder_date`` key + must still be detected (via the SystemMessage-scoped content fallback) so + in-flight conversations from before the upgrade do not re-inject. + """ + mw = _make_middleware() + state = { + "messages": [ + SystemMessage( + content="\n2026-05-08, Friday\n", + id="msg-1", + additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True}, # no reminder_date + ), + HumanMessage(content="Hello", id="msg-1__user"), + AIMessage(content="Hi there"), + 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" + result = mw.before_agent(state, _fake_runtime()) + + assert result is None # same day detected from content → no re-injection + + def test_injects_only_into_first_human_message_not_later_ones(): """Reminder targets the first HumanMessage; subsequent messages are not touched.""" mw = _make_middleware() @@ -272,10 +410,9 @@ def test_midnight_crossing_injects_date_update_as_separate_message(): """When the date has changed, a separate date-update reminder is injected before the current turn's HumanMessage using the ID-swap technique.""" mw = _make_middleware() - reminder_content = "\n2026-05-08, Friday\n" state = { "messages": [ - _reminder_msg(reminder_content, "msg-1"), + _date_reminder_msg("2026-05-08, Friday", "msg-1"), HumanMessage(content="Hello", id="msg-1__user"), AIMessage(content="Response"), HumanMessage(content="Good morning", id="msg-2"), @@ -308,10 +445,9 @@ def test_midnight_crossing_injects_date_update_as_separate_message(): def test_midnight_crossing_id_swap(): """Date-update reminder uses original ID; user message uses {id}__user.""" mw = _make_middleware() - reminder_content = "\n2026-05-08, Friday\n" state = { "messages": [ - _reminder_msg(reminder_content, "msg-1"), + _date_reminder_msg("2026-05-08, Friday", "msg-1"), HumanMessage(content="Next day message", id="msg-2"), ] } @@ -365,16 +501,12 @@ def test_memory_message_carries_reminder_key_for_title_eligibility(): def test_no_second_midnight_injection_once_date_updated(): """After a midnight update is persisted, the same-day path skips re-injection.""" mw = _make_middleware() - date_update_content = "\n2026-05-09, Saturday\n" state = { "messages": [ - _reminder_msg( - "\n2026-05-08, Friday\n", - "msg-1", - ), + _date_reminder_msg("2026-05-08, Friday", "msg-1"), HumanMessage(content="Hello", id="msg-1__user"), AIMessage(content="Response"), - _reminder_msg(date_update_content, "msg-2"), + _date_reminder_msg("2026-05-09, Saturday", "msg-2"), HumanMessage(content="Good morning", id="msg-2__user"), AIMessage(content="Good morning!"), HumanMessage(content="Third turn", id="msg-3"), diff --git a/backend/tests/test_summarization_middleware.py b/backend/tests/test_summarization_middleware.py index d01c4479a..0ab526d4f 100644 --- a/backend/tests/test_summarization_middleware.py +++ b/backend/tests/test_summarization_middleware.py @@ -7,7 +7,7 @@ from unittest.mock import MagicMock import pytest from langchain.agents import create_agent from langchain_core.language_models import BaseChatModel -from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, ToolMessage +from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, SystemMessage, ToolMessage from langchain_core.outputs import ChatGeneration, ChatResult from langgraph.constants import TAG_NOSTREAM @@ -43,11 +43,13 @@ class _StaticChatModel(BaseChatModel): return self._generate(messages, stop=stop, run_manager=run_manager, **kwargs) -def _dynamic_context_reminder(msg_id: str = "reminder-1") -> HumanMessage: - return HumanMessage( +def _dynamic_context_reminder(msg_id: str = "reminder-1") -> SystemMessage: + # Current production shape: a date SystemMessage carrying the authoritative + # date in additional_kwargs (see DynamicContextMiddleware). + return SystemMessage( content="\n2026-05-08, Friday\n", id=msg_id, - additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True}, + additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True, "reminder_date": "2026-05-08, Friday"}, )