vrr/deer-flow
2
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-07-09 16:08:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

fix(security): inject system context as SystemMessage for role isolation (#3630) (#3661)

Browse source 
P0: DynamicContextMiddleware role isolation — framework data as SystemMessage,
memory as HumanMessage (OWASP LLM01), user input unchanged.

Includes E2E hash stabilization: frontend hide_from_ui filter fix in
suggest_agent, backend _canonical_messages() hide_from_ui skip, fixture
hash updates.

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
Zhipeng Zheng 2026-06-21 20:56:42 +08:00 • committed by GitHub
parent 5b61214f7b
commit c0ce759763
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 162 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

115
backend/packages/harness/deerflow/agents/middlewares/dynamic_context_middleware.py
View file

@ -3,11 +3,11 @@
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 <system-reminder> HumanMessage inserted before the
per conversation as a dedicated <system-reminder> 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 HumanMessage before the current turn.
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.
@ -36,7 +36,7 @@ from datetime import datetime
from typing import TYPE_CHECKING, override
from langchain.agents.middleware import AgentMiddleware
from langchain_core.messages import HumanMessage
from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.runtime import Runtime
if TYPE_CHECKING:
@ -63,7 +63,10 @@ def _extract_date(content: str) -> str | 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))
# 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:
@ -86,7 +89,7 @@ def _is_user_injection_target(message: object) -> bool:
class DynamicContextMiddleware(AgentMiddleware):
"""Inject memory and current date into HumanMessages as a <system-reminder>.
"""Inject memory and current date as a SystemMessage <system-reminder>.
First turn
----------
@ -108,22 +111,31 @@ class DynamicContextMiddleware(AgentMiddleware):
self._agent_name = agent_name
self._app_config = app_config
def _build_full_reminder(self) -> str:
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
# Memory injection is gated by injection_enabled; date is always included.
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")
lines: list[str] = ["<system-reminder>"]
if memory_context:
lines.append(memory_context.strip())
lines.append("") # blank line separating memory from date
lines.append(f"<current_date>{current_date}</current_date>")
lines.append("</system-reminder>")
date_reminder = "\n".join(
[
"<system-reminder>",
f"<current_date>{current_date}</current_date>",
"</system-reminder>",
]
)
return "\n".join(lines)
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")
@ -136,29 +148,51 @@ class DynamicContextMiddleware(AgentMiddleware):
)
@staticmethod
def _make_reminder_and_user_messages(original: HumanMessage, reminder_content: str) -> tuple[HumanMessage, HumanMessage]:
"""Return (reminder_msg, user_msg) using the ID-swap technique.
def _make_reminder_and_user_messages(
original: HumanMessage,
reminder_content: str,
memory_content: str | None = None,
) -> list[SystemMessage | HumanMessage]:
"""Return messages using the ID-swap technique.
reminder_msg takes the original message's ID so that add_messages replaces it
in-place (preserving position). user_msg carries the original content with a
derived ``{id}__user`` ID and is appended immediately after by add_messages.
SystemMessage carries framework-owned data (date, metadata) — takes
the original ID so add_messages replaces it in-place. Optional
HumanMessage carries user-owned memory content with ``{id}__memory``.
The actual user message gets ``{id}__user``.
If the original message has no ID a stable UUID is generated so the derived
``{id}__user`` ID never collapses to the ambiguous ``None__user`` string.
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).
"""
stable_id = original.id or str(uuid.uuid4())
reminder_msg = HumanMessage(
content=reminder_content,
id=stable_id,
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
messages: list[SystemMessage | HumanMessage] = []
messages.append(
SystemMessage(
content=reminder_content,
id=stable_id,
additional_kwargs={"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True},
)
)
user_msg = HumanMessage(
content=original.content,
id=f"{stable_id}__user",
name=original.name,
additional_kwargs=original.additional_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 reminder_msg, user_msg
return messages
def _inject(self, state) -> dict | None:
messages = list(state.get("messages", []))
@ -175,32 +209,31 @@ class DynamicContextMiddleware(AgentMiddleware):
)
if last_date is None:
# ── First turn: inject full reminder as a separate HumanMessage ─────
# ── 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
full_reminder = self._build_full_reminder()
date_reminder, memory_block = self._build_full_reminder()
logger.info(
"DynamicContextMiddleware: injecting full reminder (len=%d, has_memory=%s) into first HumanMessage id=%r",
len(full_reminder),
"<memory>" in full_reminder,
"DynamicContextMiddleware: injecting full reminder (has_memory=%s) into first HumanMessage id=%r",
memory_block is not None,
messages[first_idx].id,
)
reminder_msg, user_msg = self._make_reminder_and_user_messages(messages[first_idx], full_reminder)
return {"messages": [reminder_msg, user_msg]}
result_msgs = self._make_reminder_and_user_messages(messages[first_idx], date_reminder, memory_block)
return {"messages": result_msgs}
if last_date == current_date:
# ── Same day: nothing to do ──────────────────────────────────────────
return None
# ── Midnight crossed: inject date-update reminder as a separate HumanMessage ──
# ── 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
reminder_msg, user_msg = 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())
logger.info("DynamicContextMiddleware: midnight crossing detected — injected date update before current turn")
return {"messages": [reminder_msg, user_msg]}
return {"messages": result_msgs}
@override
def before_agent(self, state, runtime: Runtime) -> dict | None:

8
backend/tests/fixtures/replay/write_read_file.ultra.json vendored
View file

@ -59,8 +59,8 @@
},
{
"caller": "middleware:title",
"conversation_hash": "3598aeb87e221ca8f554e4d61ce6d5e8801754606fa5c95a89c38bd6cb623045",
"input_hash": "75101f9faa453b1a35deff920b1e3c1a9f0b013a7627fbbaa03436752776b953",
"conversation_hash": "958df7bac4e74830d46d29bc19cdea85cae2fc3ebc65ca00a9a5d308bba8ad30",
"input_hash": "f1d96cb1f7cfbe6baf6b2c7b82d32f5561656b1c1492e2714c51631808b50b79",
"output": {
"type": "ai",
"data": {
@ -174,8 +174,8 @@
},
{
"caller": "suggest_agent",
"conversation_hash": "8b98ebdbb53e88f000556c4753adede8eaa076ff6fd7b8a1285bfd18aee8144d",
"input_hash": "dcd855d389d7179a1e4bc7074fa9ba7ce697570af8947225d6bacb538f14a0cb",
"conversation_hash": "036fcf61f733a11153586db4fc8c39c412c90ac99d5043c22e812c90572176bf",
"input_hash": "d8b5968661b908283a8c3aadae3f4529ba3a5b0098849b0dde330d5b6ce7941a",
"output": {
"type": "ai",
"data": {

22
backend/tests/replay_provider.py
View file

@ -156,8 +156,17 @@ _BOUNDARY_BEGIN_RE = re.compile(r"--- BEGIN USER INPUT ---\n?")
_BOUNDARY_END_RE = re.compile(r"\n?--- END USER INPUT ---")
# After _SYSTEM_REMINDER_RE strips a <system-reminder> block, a role label like
# "User: " or "Assistant: " may remain with nothing after it (just whitespace/newline).
# Collapse those empty role lines so the hash is resilient to reminder leaks from the
# frontend (e.g. a HumanMessage(hide_from_ui=True) whose <system-reminder> content
# was stripped, leaving "User: \n" residue that would otherwise cause a hash mismatch).
_EMPTY_ROLE_LINE_RE = re.compile(r"^(User|Assistant):\s*$\n?", re.MULTILINE)
def _normalize_text(text: str) -> str:
text = _SYSTEM_REMINDER_RE.sub("", text)
text = _EMPTY_ROLE_LINE_RE.sub("", text)
text = _BOUNDARY_BEGIN_RE.sub("", text)
text = _BOUNDARY_END_RE.sub("", text)
text = _UUID_RE.sub("<UUID>", text)
@ -202,6 +211,19 @@ def _canonical_messages(messages: list[BaseMessage]) -> str:
# most-edited part of the prompt and not part of the contract under test.
if message.type == "system":
continue
# Exclude framework-injected hidden messages (dynamic-context reminders,
# memory injections, etc.) regardless of type. These carry volatile
# per-session data (current date, user memory) and are not user-authored
# content, so they must not participate in the match key. On the p1
# branch they were HumanMessages whose <system-reminder> content was
# stripped → empty → excluded by the empty-content check below. On p0
# they may be SystemMessages (excluded by type) or HumanMessages with
# standalone <memory> content (not stripped by _SYSTEM_REMINDER_RE,
# not empty, so previously leaked into the hash). Checking hide_from_ui
# directly makes the hash stable across all middleware implementations.
additional_kwargs = getattr(message, "additional_kwargs", None) or {}
if additional_kwargs.get("hide_from_ui"):
continue
content = _normalize_text(_content_to_text(message.content))
tool_calls = getattr(message, "tool_calls", None)
# Drop messages that are empty after normalization — e.g. a turn that was

68
backend/tests/test_dynamic_context_middleware.py
View file

@ -7,7 +7,7 @@ the first HumanMessage exactly once per session (frozen-snapshot pattern).
from types import SimpleNamespace
from unittest import mock
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from deerflow.agents.middlewares.dynamic_context_middleware import (
_DYNAMIC_CONTEXT_REMINDER_KEY,
@ -26,7 +26,11 @@ def _fake_runtime():
def _reminder_msg(content: str, msg_id: str) -> HumanMessage:
"""Build a reminder HumanMessage the way the middleware would produce it."""
"""Build a pre-PR HumanMessage reminder — simulates historical checkpoints.
Uses HumanMessage (DEPRECATED format) to exercise the backward-compat
path in ``is_dynamic_context_reminder``. New reminders are SystemMessage.
"""
return HumanMessage(
content=content,
id=msg_id,
@ -52,7 +56,7 @@ def test_injects_system_reminder_into_first_human_message():
assert len(updated_msgs) == 2
reminder_msg = updated_msgs[0]
assert isinstance(reminder_msg, HumanMessage)
assert isinstance(reminder_msg, SystemMessage)
assert reminder_msg.id == "msg-1" # takes the original ID (position swap)
assert reminder_msg.additional_kwargs.get(_DYNAMIC_CONTEXT_REMINDER_KEY) is True
assert _SYSTEM_REMINDER_TAG in reminder_msg.content
@ -79,11 +83,18 @@ def test_memory_included_when_present():
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
result = mw.before_agent(state, _fake_runtime())
# Reminder is the first returned message; user query is the second
reminder_content = result["messages"][0].content
assert "User prefers Python." in reminder_content
assert "<current_date>2026-05-08, Friday</current_date>" in reminder_content
assert result["messages"][1].content == "Hi"
# Memory is a separate HumanMessage — not merged into SystemMessage (OWASP LLM01)
msgs = result["messages"]
assert len(msgs) == 3 # date SystemMessage + memory HumanMessage + user HumanMessage
assert isinstance(msgs[0], SystemMessage)
assert "<current_date>2026-05-08, Friday</current_date>" in msgs[0].content
assert "User prefers Python." not in msgs[0].content # memory NOT in system role
assert isinstance(msgs[1], HumanMessage)
assert "User prefers Python." in msgs[1].content
assert msgs[2].content == "Hi"
# ---------------------------------------------------------------------------
@ -279,6 +290,9 @@ def test_midnight_crossing_injects_date_update_as_separate_message():
msgs = result["messages"]
assert len(msgs) == 2
# Midnight-cross reminder is also a SystemMessage — both paths are covered
assert isinstance(msgs[0], SystemMessage)
# Date-update reminder takes the current message's ID
assert msgs[0].id == "msg-2"
assert msgs[0].additional_kwargs.get(_DYNAMIC_CONTEXT_REMINDER_KEY) is True
@ -310,6 +324,44 @@ def test_midnight_crossing_id_swap():
assert result["messages"][1].id == "msg-2__user"
def test_memory_message_carries_reminder_key_for_title_eligibility():
"""Regression: memory HumanMessage must carry _DYNAMIC_CONTEXT_REMINDER_KEY.
Without it, title_middleware._is_user_message_for_title counts the memory
block as a second user message and skips title generation entirely.
Similarly, summarization_middleware._preserve_dynamic_context_reminders
would not rescue the memory block from summary compression.
"""
from deerflow.agents.middlewares.dynamic_context_middleware import is_dynamic_context_reminder
mw = _make_middleware()
state = {"messages": [HumanMessage(content="Hi", id="msg-1")]}
with (
mock.patch(
"deerflow.agents.lead_agent.prompt._get_memory_context",
return_value="<memory>\nUser prefers Python.\n</memory>",
),
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"]
# Memory message must be recognized as a dynamic-context reminder
memory_msg = msgs[1]
assert isinstance(memory_msg, HumanMessage)
assert memory_msg.id == "msg-1__memory"
assert is_dynamic_context_reminder(memory_msg) is True
# Only the actual user message is title-eligible
from deerflow.agents.middlewares.title_middleware import TitleMiddleware
title_eligible = [m for m in msgs if TitleMiddleware._is_user_message_for_title(m)]
assert len(title_eligible) == 1
assert title_eligible[0].content == "Hi"
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()

2
frontend/src/components/workspace/input-box.tsx
View file

@ -59,6 +59,7 @@ import {
import { fetch } from "@/core/api/fetcher";
import { getBackendBaseURL } from "@/core/config";
import { useI18n } from "@/core/i18n/hooks";
import { isHiddenFromUIMessage } from "@/core/messages/utils";
import { useModels } from "@/core/models/hooks";
import type { Skill } from "@/core/skills";
import { useSkills } from "@/core/skills/hooks";
@ -533,6 +534,7 @@ export function InputBox({
const recent = messagesRef.current
.filter((m) => m.type === "human" || m.type === "ai")
.filter((m) => !isHiddenFromUIMessage(m))
.map((m) => {
const role = m.type === "human" ? "user" : "assistant";
const content = textOfMessage(m) ?? "";