"""Middleware to inject dynamic context (memory, current date) as a system-reminder. The system prompt is kept fully static for maximum prefix-cache reuse across users and sessions. The current date is always injected. Per-user memory is also injected when ``memory.injection_enabled`` is True in the app config. Both are delivered once per conversation as a dedicated SystemMessage inserted before the first user message (frozen-snapshot pattern). When a conversation spans midnight the middleware detects the date change and injects a lightweight date-update reminder as a separate SystemMessage before the current turn. This correction is persisted so subsequent turns on the new day see a consistent history and do not re-inject. Reminder format: ... 2026-05-08, Friday Date-update format: 2026-05-09, Saturday """ from __future__ import annotations import asyncio import logging import re import uuid from datetime import datetime from typing import TYPE_CHECKING, override from langchain.agents.middleware import AgentMiddleware from langchain_core.messages import HumanMessage, SystemMessage from langgraph.runtime import Runtime if TYPE_CHECKING: from deerflow.config.app_config import AppConfig logger = logging.getLogger(__name__) # Upper bound (seconds) for a single _inject() offload. If the warm-up at # gateway startup failed silently, the first request may still hit a cold # tiktoken BPE download that blocks until the OS TCP timeout (~26 min). # This cap ensures the request degrades gracefully instead of hanging. _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" def _extract_date(content: str) -> str | None: """Return the first value found in *content*, or None.""" m = _DATE_RE.search(content) return m.group(1) if m else None def is_dynamic_context_reminder(message: object) -> bool: """Return whether *message* is a hidden dynamic-context reminder.""" # DEPRECATED: HumanMessage reminders only exist in pre-PR checkpoints. # Once all active checkpoints are migrated, the HumanMessage branch can be # removed and this function can check SystemMessage exclusively. return isinstance(message, (HumanMessage, SystemMessage)) and bool(message.additional_kwargs.get(_DYNAMIC_CONTEXT_REMINDER_KEY)) def _last_injected_date(messages: list) -> str | None: """Scan messages in reverse and return the most recently injected date. 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 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) date = _extract_date(content_str) if date is not None: return date return None def _is_user_injection_target(message: object) -> bool: """Return whether *message* can receive a dynamic-context reminder.""" 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): """Inject memory and current date as a SystemMessage . First turn ---------- Prepends a full system-reminder (memory + date) to the first HumanMessage and persists it (same message ID). The first message is then frozen for the whole session — its content never changes again, so the prefix cache can hit on every subsequent turn. Midnight crossing ----------------- If the conversation spans midnight, the current date differs from the date that was injected earlier. In that case a lightweight date-update reminder is prepended to the **current** (last) HumanMessage and persisted. Subsequent turns on the new day see the corrected date in history and skip re-injection. """ def __init__(self, agent_name: str | None = None, *, app_config: AppConfig | None = None): super().__init__() self._agent_name = agent_name self._app_config = app_config def _build_full_reminder(self) -> tuple[str, str | None]: """Return (date_reminder, memory_block | None). Framework-owned data (date) is separated from user-owned data (memory) so the downstream SystemMessage carries only framework authority and memory stays at role:user — preventing untrusted content from gaining system privilege (OWASP LLM01). """ from deerflow.agents.lead_agent.prompt import _get_memory_context injection_enabled = self._app_config.memory.injection_enabled if self._app_config else True memory_context = _get_memory_context(self._agent_name, app_config=self._app_config) if injection_enabled else "" current_date = datetime.now().strftime("%Y-%m-%d, %A") date_reminder = "\n".join( [ "", f"{current_date}", "", ] ) memory_block = memory_context.strip() if memory_context else None return date_reminder, memory_block def _build_date_update_reminder(self) -> str: current_date = datetime.now().strftime("%Y-%m-%d, %A") return "\n".join( [ "", f"{current_date}", "", ] ) @staticmethod def _make_reminder_and_user_messages( 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. *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) — 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=reminder_kwargs, ) ) if memory_content: messages.append( HumanMessage( content=memory_content, id=f"{stable_id}__memory", additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True}, ) ) messages.append( HumanMessage( content=original.content, id=f"{stable_id}__user", name=original.name, additional_kwargs=original.additional_kwargs, ) ) return messages def _inject(self, state) -> dict | None: messages = list(state.get("messages", [])) if not messages: return None current_date = datetime.now().strftime("%Y-%m-%d, %A") last_date = _last_injected_date(messages) logger.debug( "DynamicContextMiddleware._inject: msg_count=%d last_date=%r current_date=%r", len(messages), last_date, current_date, ) if last_date is None: # ── First turn: inject full reminder as a SystemMessage ───── first_idx = next((i for i, m in enumerate(messages) if _is_user_injection_target(m)), None) if first_idx is None: return None date_reminder, memory_block = self._build_full_reminder() logger.info( "DynamicContextMiddleware: injecting full reminder (has_memory=%s) into first HumanMessage id=%r", memory_block is not None, messages[first_idx].id, ) 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: # ── Same day: nothing to do ────────────────────────────────────────── return None # ── Midnight crossed: inject date-update reminder as a SystemMessage ── 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 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} @override def before_agent(self, state, runtime: Runtime) -> dict | None: return self._inject(state) @override async def abefore_agent(self, state, runtime: Runtime) -> dict | None: # _inject() performs synchronous file I/O (memory JSON loading) and # potentially blocking network calls (tiktoken encoding download on # first use). Offload to a thread so the event loop is never blocked # — a blocking call here starves all concurrent HTTP handlers (auth, # SSE heartbeats, etc.). See issue #3402. # # Bounded timeout: if startup warm-up failed silently (e.g. network # blip during deploy), the first request's cold tiktoken download can # block for tens of minutes (OS TCP timeout). Time-box injection so # the request degrades gracefully (no memory context) rather than # hanging. try: return await asyncio.wait_for( asyncio.to_thread(self._inject, state), timeout=_INJECT_TIMEOUT_SECONDS, ) except TimeoutError: logger.warning( "DynamicContextMiddleware: injection timed out (%.1fs); skipping memory/date injection for this turn", _INJECT_TIMEOUT_SECONDS, ) return None