mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-09 16:08:31 +00:00
fix(agents): coalesce SystemMessages before LLM request (#3711)
Some checks are pending
Backend Blocking IO / backend-blocking-io (push) Waiting to run
Unit Tests / backend-unit-tests (push) Waiting to run
Frontend Unit Tests / frontend-unit-tests (push) Waiting to run
Lint Check / lint-backend (push) Waiting to run
Lint Check / lint-frontend (push) Waiting to run
Replay E2E (front-back contract) / Layer 1 — backend golden (no API key) (push) Waiting to run
Replay E2E (front-back contract) / Layer 2 — full-stack render (no API key) (push) Waiting to run
Some checks are pending
Backend Blocking IO / backend-blocking-io (push) Waiting to run
Unit Tests / backend-unit-tests (push) Waiting to run
Frontend Unit Tests / frontend-unit-tests (push) Waiting to run
Lint Check / lint-backend (push) Waiting to run
Lint Check / lint-frontend (push) Waiting to run
Replay E2E (front-back contract) / Layer 1 — backend golden (no API key) (push) Waiting to run
Replay E2E (front-back contract) / Layer 2 — full-stack render (no API key) (push) Waiting to run
DynamicContextMiddleware (PR #3630, p0) uses the ID-swap technique to inject a SystemMessage(reminder) into the middle of the conversation. create_agent then prepends the static system_prompt as another SystemMessage at request time, so strict OpenAI-compatible backends (vLLM, SGLang, Qwen) and Anthropic reject the request with 'System message must be at the beginning'. New SystemMessageCoalescingMiddleware runs in wrap_model_call — after create_agent prepends system_prompt — and merges every SystemMessage into a single leading one before the request reaches the provider. Non-system messages keep their original order; the merged SystemMessage preserves the id of the first system message. Only the request payload is touched; checkpoint state is unchanged, so every consumer that scans history (memory builder, journal, summarization, dynamic-context detection) keeps working. Mirrors the per-request coalescing already done for Claude in claude_provider._coalesce_system_messages (PR #3702) but at a provider-agnostic layer so every backend benefits from a single fix. Closes #3707
This commit is contained in:
parent
cefc53c72a
commit
fde6885aea
4 changed files with 843 additions and 4 deletions
|
|
@ -30,7 +30,7 @@ deer-flow/
|
|||
│ │ └── deerflow/
|
||||
│ │ ├── agents/ # LangGraph agent system
|
||||
│ │ │ ├── lead_agent/ # Main agent (factory + system prompt)
|
||||
│ │ │ ├── middlewares/ # 10 middleware components
|
||||
│ │ │ ├── middlewares/ # 11 middleware components
|
||||
│ │ │ ├── memory/ # Memory extraction, queue, prompts
|
||||
│ │ │ └── thread_state.py # ThreadState schema
|
||||
│ │ ├── sandbox/ # Sandbox execution system
|
||||
|
|
@ -218,9 +218,10 @@ Lead-agent middlewares are assembled in strict append order across `packages/har
|
|||
14. **MemoryMiddleware** - Queues conversations for async memory update (filters to user + final AI responses)
|
||||
15. **ViewImageMiddleware** - Injects base64 image data before LLM call (conditional on vision support)
|
||||
16. **DeferredToolFilterMiddleware** - Hides deferred (MCP) tool schemas from the bound model using a build-time deferred-name set + catalog hash, reading per-thread promotions from `ThreadState.promoted` (hash-scoped, no ContextVar); a tool becomes bound on subsequent turns after `tool_search` returns its schema (optional, if `tool_search.enabled`)
|
||||
17. **SubagentLimitMiddleware** - Truncates excess `task` tool calls from model response to enforce `MAX_CONCURRENT_SUBAGENTS` limit (optional, if `subagent_enabled`)
|
||||
18. **LoopDetectionMiddleware** - Detects repeated tool-call loops; hard-stop responses clear both structured `tool_calls` and raw provider tool-call metadata before forcing a final text answer
|
||||
19. **ClarificationMiddleware** - Intercepts `ask_clarification` tool calls, interrupts via `Command(goto=END)` (must be last)
|
||||
17. **SystemMessageCoalescingMiddleware** - Merges `request.system_message` and every in-`request.messages` SystemMessage into a single leading SystemMessage in `wrap_model_call`; provider-agnostic fix for strict backends (vLLM/SGLang/Qwen/Anthropic) that reject non-leading system messages. Touches the per-request payload only — checkpoint state is unchanged so `is_dynamic_context_reminder` history scanners keep working. On midnight crossings, earlier `dynamic_context_reminder`-tagged SystemMessages are dropped and only the latest (most recent date) survives, avoiding two contradictory `<current_date>` blocks appearing adjacent with no temporal anchor.
|
||||
18. **SubagentLimitMiddleware** - Truncates excess `task` tool calls from model response to enforce `MAX_CONCURRENT_SUBAGENTS` limit (optional, if `subagent_enabled`)
|
||||
19. **LoopDetectionMiddleware** - Detects repeated tool-call loops; hard-stop responses clear both structured `tool_calls` and raw provider tool-call metadata before forcing a final text answer
|
||||
20. **ClarificationMiddleware** - Intercepts `ask_clarification` tool calls, interrupts via `Command(goto=END)` (must be last)
|
||||
|
||||
### Configuration System
|
||||
|
||||
|
|
|
|||
|
|
@ -348,6 +348,13 @@ def build_middlewares(
|
|||
|
||||
middlewares.append(DeferredToolFilterMiddleware(deferred_setup.deferred_names, deferred_setup.catalog_hash))
|
||||
|
||||
# Coalesce every SystemMessage into a single leading one before the request
|
||||
# reaches the provider. Strict backends (vLLM, SGLang, Qwen, Anthropic)
|
||||
# reject non-leading SystemMessages. See system_message_coalescing_middleware.py.
|
||||
from deerflow.agents.middlewares.system_message_coalescing_middleware import SystemMessageCoalescingMiddleware
|
||||
|
||||
middlewares.append(SystemMessageCoalescingMiddleware())
|
||||
|
||||
# Add SubagentLimitMiddleware to truncate excess parallel task calls
|
||||
subagent_enabled = cfg.get("subagent_enabled", False)
|
||||
if subagent_enabled:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,150 @@
|
|||
"""Middleware to coalesce multiple SystemMessages into a single leading one.
|
||||
|
||||
Strict OpenAI-compatible backends (vLLM, SGLang, Qwen) and Anthropic reject
|
||||
non-leading SystemMessages with errors like "System message must be at the
|
||||
beginning" or "Received multiple non-consecutive system messages". The
|
||||
official OpenAI API tolerates mid-conversation system messages, so the issue
|
||||
only surfaces on strict backends.
|
||||
|
||||
DeerFlow's lead agent accumulates multiple SystemMessages because
|
||||
DynamicContextMiddleware uses the ID-swap technique to replace the first or
|
||||
last HumanMessage with a triplet whose first element is a SystemMessage
|
||||
reminder (framework-owned date/metadata must not masquerade as user input,
|
||||
per OWASP LLM01). On midnight crossings a second SystemMessage (date update)
|
||||
is injected. create_agent holds the static system_prompt in the separate
|
||||
``request.system_message`` field and only flattens it into the message list
|
||||
inside the model-call handler (``[request.system_message, *messages]``).
|
||||
|
||||
This middleware runs in wrap_model_call — before the handler flattens the two
|
||||
— and merges ``request.system_message`` plus every SystemMessage found in
|
||||
``request.messages`` into a single leading SystemMessage emitted via the
|
||||
``system_message`` field. It only touches the request payload; the persistent
|
||||
conversation state (checkpoint) is unchanged, so middleware that scans history
|
||||
by marker (e.g. is_dynamic_context_reminder) keeps working.
|
||||
|
||||
Note: Mirrors the per-request coalescing already done for Claude in
|
||||
claude_provider._coalesce_system_messages but at a provider-agnostic layer so
|
||||
every backend benefits from a single fix instead of per-provider patches.
|
||||
"""
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import override
|
||||
|
||||
from langchain.agents import AgentState
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
|
||||
from langchain_core.messages import SystemMessage
|
||||
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import is_dynamic_context_reminder
|
||||
|
||||
|
||||
def _flatten_content(content) -> str:
|
||||
"""Convert message content to a plain string, handling both str and list types.
|
||||
|
||||
langchain messages support list-type content for multimodal (e.g.
|
||||
``[{"type": "text", "text": "..."}]``). SystemMessages in DeerFlow are always
|
||||
plain strings, but this helper ensures robustness for any content shape.
|
||||
"""
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts = []
|
||||
for item in content:
|
||||
if isinstance(item, str):
|
||||
parts.append(item)
|
||||
elif isinstance(item, dict) and "text" in item:
|
||||
parts.append(item["text"])
|
||||
else:
|
||||
parts.append(str(item))
|
||||
return "\n".join(parts)
|
||||
return str(content)
|
||||
|
||||
|
||||
def _coalesce_request(request: ModelRequest) -> ModelRequest | None:
|
||||
"""Merge ``request.system_message`` and in-``messages`` SystemMessages into one.
|
||||
|
||||
On langchain >= 1.2.15 the static system prompt lives in the separate
|
||||
``request.system_message`` field, not in ``request.messages``. The model-call
|
||||
handler flattens them at the very last moment (``[system_message, *messages]``),
|
||||
so a middleware that only scans ``messages`` cannot see the prompt and ends up
|
||||
a no-op. This helper inspects both sources, merges every SystemMessage into a
|
||||
single entry, and emits the result via ``system_message`` so the handler still
|
||||
prepends it correctly.
|
||||
|
||||
Returns None when no SystemMessages live inside ``messages`` — in that case
|
||||
``system_message`` (if set) is already the sole leading system block and the
|
||||
request can pass through with zero mutation, preserving prefix-cache hits.
|
||||
"""
|
||||
in_msg_systems = [m for m in request.messages if isinstance(m, SystemMessage)]
|
||||
if not in_msg_systems:
|
||||
return None
|
||||
|
||||
# Merge system_message (if any) + all in-messages SystemMessages.
|
||||
parts: list[SystemMessage] = []
|
||||
if request.system_message is not None:
|
||||
parts.append(request.system_message)
|
||||
parts.extend(in_msg_systems)
|
||||
|
||||
# Deduplicate dynamic_context_reminder SystemMessages: only keep the last
|
||||
# one (most recent date), drop earlier reminders. On midnight crossings
|
||||
# the merged content would otherwise contain two adjacent contradictory
|
||||
# <current_date> blocks with no temporal anchor — the intervening turns
|
||||
# that originally separated them are gone after coalescing. The model
|
||||
# should see only the latest date, not a stale one it must guess to
|
||||
# ignore.
|
||||
reminder_indices = [i for i, p in enumerate(parts) if is_dynamic_context_reminder(p)]
|
||||
if len(reminder_indices) > 1:
|
||||
keep_last = reminder_indices[-1]
|
||||
parts = [p for i, p in enumerate(parts) if i not in reminder_indices[:-1] or i == keep_last]
|
||||
|
||||
# Preserve the id of the first SystemMessage (typically the static
|
||||
# system_prompt) so downstream consumers that key off the leading system
|
||||
# message id are unaffected. Merge additional_kwargs from all parts so
|
||||
# markers like hide_from_ui / dynamic_context_reminder from reminders are
|
||||
# retained on the coalesced block.
|
||||
first = parts[0]
|
||||
merged_kwargs: dict = {}
|
||||
for p in parts:
|
||||
merged_kwargs.update(p.additional_kwargs or {})
|
||||
merged = SystemMessage(
|
||||
content="\n\n".join(_flatten_content(p.content) for p in parts),
|
||||
id=first.id,
|
||||
additional_kwargs=merged_kwargs,
|
||||
)
|
||||
|
||||
non_system = [m for m in request.messages if not isinstance(m, SystemMessage)]
|
||||
return request.override(system_message=merged, messages=non_system)
|
||||
|
||||
|
||||
class SystemMessageCoalescingMiddleware(AgentMiddleware[AgentState]):
|
||||
"""Merge all SystemMessages into a single leading SystemMessage.
|
||||
|
||||
Uses wrap_model_call (not before_agent) so the merge runs on the final
|
||||
request payload — where ``system_message`` and ``messages`` are still
|
||||
separate fields — and never touches the persisted state["messages"]. This
|
||||
keeps the checkpoint structure intact for every consumer that scans history
|
||||
(memory builder, journal, summarization, dynamic-context detection).
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _maybe_coalesce(request: ModelRequest) -> ModelRequest:
|
||||
coalesced = _coalesce_request(request)
|
||||
if coalesced is None:
|
||||
return request
|
||||
return coalesced
|
||||
|
||||
@override
|
||||
def wrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], ModelResponse],
|
||||
) -> ModelCallResult:
|
||||
return handler(self._maybe_coalesce(request))
|
||||
|
||||
@override
|
||||
async def awrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
||||
) -> ModelCallResult:
|
||||
return await handler(self._maybe_coalesce(request))
|
||||
681
backend/tests/test_system_message_coalescing_middleware.py
Normal file
681
backend/tests/test_system_message_coalescing_middleware.py
Normal file
|
|
@ -0,0 +1,681 @@
|
|||
"""Tests for SystemMessageCoalescingMiddleware.
|
||||
|
||||
Verifies that ``request.system_message`` and in-``messages`` SystemMessages are
|
||||
merged into a single leading SystemMessage before the request reaches the LLM,
|
||||
fixing the "System message must be at the beginning" error on strict
|
||||
OpenAI-compatible backends (vLLM, SGLang, Qwen) and Anthropic.
|
||||
|
||||
On langchain >= 1.2.15 the static system prompt lives in the separate
|
||||
``request.system_message`` field, not in ``request.messages``. The model-call
|
||||
handler flattens them at the very last moment (``[system_message, *messages]``),
|
||||
so tests must build requests with the same split to exercise the real code path.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
|
||||
|
||||
from deerflow.agents.middlewares.system_message_coalescing_middleware import (
|
||||
SystemMessageCoalescingMiddleware,
|
||||
_coalesce_request,
|
||||
_flatten_content,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_request(system_message: SystemMessage | None, messages: list[BaseMessage]):
|
||||
"""Build a minimal ModelRequest stand-in matching langchain 1.2.15 shape.
|
||||
|
||||
``system_message`` and ``messages`` are separate fields — the handler
|
||||
flattens them into ``[system_message, *messages]`` only at call time.
|
||||
"""
|
||||
request = MagicMock()
|
||||
request.system_message = system_message
|
||||
request.messages = list(messages)
|
||||
request.override = lambda **updates: _override_request(request, updates)
|
||||
return request
|
||||
|
||||
|
||||
def _override_request(request, updates):
|
||||
"""Mimic ModelRequest.override(): return a copy with fields replaced."""
|
||||
new = MagicMock()
|
||||
new.system_message = updates.get("system_message", request.system_message)
|
||||
new.messages = updates.get("messages", request.messages)
|
||||
new.override = lambda **kw: _override_request(new, kw)
|
||||
return new
|
||||
|
||||
|
||||
def _capture_handler():
|
||||
"""Return (captured_requests, handler) that records what was sent."""
|
||||
captured = []
|
||||
|
||||
def handler(req):
|
||||
captured.append(req)
|
||||
return "response"
|
||||
|
||||
return captured, handler
|
||||
|
||||
|
||||
def _final_payload(request) -> list[BaseMessage]:
|
||||
"""Simulate what the LLM receives: [system_message, *messages] (if set)."""
|
||||
if request.system_message is not None:
|
||||
return [request.system_message, *request.messages]
|
||||
return list(request.messages)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _coalesce_request (pure helper)
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestCoalesceRequest:
|
||||
def test_no_system_anywhere_returns_none(self):
|
||||
"""Zero SystemMessages → passthrough."""
|
||||
request = _make_request(system_message=None, messages=[HumanMessage(content="hi")])
|
||||
assert _coalesce_request(request) is None
|
||||
|
||||
def test_only_system_message_field_returns_none(self):
|
||||
"""system_message set, no SystemMessages in messages → passthrough.
|
||||
|
||||
The single system block is already in the right place; coalescing would
|
||||
just create a new object with the same content (zero drift).
|
||||
"""
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="prompt"),
|
||||
messages=[HumanMessage(content="hi")],
|
||||
)
|
||||
assert _coalesce_request(request) is None
|
||||
|
||||
def test_system_message_plus_one_in_msg_system_coalesces(self):
|
||||
"""system_message + 1 in-messages SystemMessage → merge into system_message."""
|
||||
prompt = SystemMessage(content="You are DeerFlow.", id="sys-1")
|
||||
reminder = SystemMessage(content="<system-reminder>date</system-reminder>", id="msg-1")
|
||||
user = HumanMessage(content="Hello", id="msg-1__user")
|
||||
request = _make_request(system_message=prompt, messages=[reminder, user])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert result.system_message is not None
|
||||
assert "You are DeerFlow." in result.system_message.content
|
||||
assert "<system-reminder>date</system-reminder>" in result.system_message.content
|
||||
# messages no longer contain any SystemMessage
|
||||
assert not any(isinstance(m, SystemMessage) for m in result.messages)
|
||||
# user message preserved
|
||||
assert any(m.content == "Hello" for m in result.messages)
|
||||
|
||||
def test_system_message_plus_two_in_msg_systems_coalesces(self):
|
||||
"""system_message + 2 non-reminder SystemMessages → merge all (no dedup)."""
|
||||
prompt = SystemMessage(content="prompt")
|
||||
reminder = SystemMessage(content="<system-reminder>day1</system-reminder>")
|
||||
date_update = SystemMessage(content="<system-reminder>day2</system-reminder>")
|
||||
user = HumanMessage(content="next")
|
||||
request = _make_request(system_message=prompt, messages=[reminder, date_update, user])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert "prompt" in result.system_message.content
|
||||
assert "day1" in result.system_message.content
|
||||
assert "day2" in result.system_message.content
|
||||
|
||||
def test_no_system_message_but_in_msg_systems_coalesces(self):
|
||||
"""system_message is None, but messages has SystemMessages → move to system_message."""
|
||||
reminder = SystemMessage(content="reminder")
|
||||
user = HumanMessage(content="hi")
|
||||
request = _make_request(system_message=None, messages=[reminder, user])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert result.system_message is not None
|
||||
assert "reminder" in result.system_message.content
|
||||
assert not any(isinstance(m, SystemMessage) for m in result.messages)
|
||||
|
||||
def test_merged_content_uses_double_newline_separator(self):
|
||||
"""System contents are joined with \\n\\n."""
|
||||
prompt = SystemMessage(content="PART_A")
|
||||
reminder = SystemMessage(content="PART_B")
|
||||
request = _make_request(system_message=prompt, messages=[reminder])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert result.system_message.content == "PART_A\n\nPART_B"
|
||||
|
||||
def test_merged_preserves_first_system_message_id(self):
|
||||
"""The merged SystemMessage keeps the id of system_message (first in order)."""
|
||||
prompt = SystemMessage(content="prompt", id="sys-1")
|
||||
reminder = SystemMessage(content="reminder", id="msg-1")
|
||||
request = _make_request(system_message=prompt, messages=[reminder])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert result.system_message.id == "sys-1"
|
||||
|
||||
def test_merged_preserves_in_msg_id_when_no_system_message(self):
|
||||
"""When system_message is None, merged id comes from the first in-msg system."""
|
||||
reminder = SystemMessage(content="reminder", id="msg-1")
|
||||
request = _make_request(system_message=None, messages=[reminder])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert result.system_message.id == "msg-1"
|
||||
|
||||
def test_non_system_messages_keep_original_order(self):
|
||||
"""HumanMessage/AIMessage order is preserved after coalescing."""
|
||||
prompt = SystemMessage(content="prompt")
|
||||
user1 = HumanMessage(content="u1", id="u1")
|
||||
ai = AIMessage(content="a1", id="a1")
|
||||
reminder = SystemMessage(content="reminder")
|
||||
user2 = HumanMessage(content="u2", id="u2")
|
||||
request = _make_request(system_message=prompt, messages=[user1, ai, reminder, user2])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
non_system = result.messages
|
||||
assert [m.id for m in non_system] == ["u1", "a1", "u2"]
|
||||
|
||||
def test_merged_kwargs_combine_all_parts(self):
|
||||
"""additional_kwargs from all parts are merged into the result."""
|
||||
prompt = SystemMessage(
|
||||
content="prompt",
|
||||
id="sys-1",
|
||||
additional_kwargs={"source": "prompt"},
|
||||
)
|
||||
reminder = SystemMessage(
|
||||
content="reminder",
|
||||
id="msg-1",
|
||||
additional_kwargs={"hide_from_ui": True, "dynamic_context_reminder": True},
|
||||
)
|
||||
request = _make_request(system_message=prompt, messages=[reminder])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert result.system_message.additional_kwargs == {
|
||||
"source": "prompt",
|
||||
"hide_from_ui": True,
|
||||
"dynamic_context_reminder": True,
|
||||
}
|
||||
|
||||
def test_merged_kwargs_later_parts_override(self):
|
||||
"""When two parts share a key, the later part's value wins."""
|
||||
prompt = SystemMessage(
|
||||
content="prompt",
|
||||
id="sys-1",
|
||||
additional_kwargs={"priority": "low"},
|
||||
)
|
||||
reminder = SystemMessage(
|
||||
content="reminder",
|
||||
id="msg-1",
|
||||
additional_kwargs={"priority": "high"},
|
||||
)
|
||||
request = _make_request(system_message=prompt, messages=[reminder])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert result.system_message.additional_kwargs["priority"] == "high"
|
||||
|
||||
def test_merged_handles_list_content(self):
|
||||
"""List-type SystemMessage content is flattened before joining."""
|
||||
prompt = SystemMessage(
|
||||
content=[{"type": "text", "text": "You are DeerFlow."}],
|
||||
id="sys-1",
|
||||
)
|
||||
reminder = SystemMessage(content="<system-reminder>date</system-reminder>", id="msg-1")
|
||||
request = _make_request(system_message=prompt, messages=[reminder])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert "You are DeerFlow." in result.system_message.content
|
||||
assert "<system-reminder>date</system-reminder>" in result.system_message.content
|
||||
|
||||
def test_reminder_dedup_keeps_only_last(self):
|
||||
"""When multiple SystemMessages have dynamic_context_reminder=True,
|
||||
only the last one survives; earlier ones are dropped."""
|
||||
prompt = SystemMessage(content="prompt", id="sys-prompt")
|
||||
day1 = SystemMessage(
|
||||
content="<system-reminder>day1</system-reminder>",
|
||||
id="msg-1",
|
||||
additional_kwargs={"hide_from_ui": True, "dynamic_context_reminder": True},
|
||||
)
|
||||
day2 = SystemMessage(
|
||||
content="<system-reminder>day2</system-reminder>",
|
||||
id="msg-2",
|
||||
additional_kwargs={"hide_from_ui": True, "dynamic_context_reminder": True},
|
||||
)
|
||||
user = HumanMessage(content="hi")
|
||||
request = _make_request(system_message=prompt, messages=[day1, day2, user])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert "prompt" in result.system_message.content
|
||||
assert "day2" in result.system_message.content
|
||||
assert "day1" not in result.system_message.content
|
||||
|
||||
def test_reminder_dedup_does_not_affect_non_reminder_systems(self):
|
||||
"""SystemMessages without dynamic_context_reminder are never deduplicated."""
|
||||
prompt = SystemMessage(content="prompt", id="sys-prompt")
|
||||
other = SystemMessage(content="custom system block", id="msg-1")
|
||||
reminder = SystemMessage(
|
||||
content="<system-reminder>day2</system-reminder>",
|
||||
id="msg-2",
|
||||
additional_kwargs={"dynamic_context_reminder": True},
|
||||
)
|
||||
user = HumanMessage(content="hi")
|
||||
request = _make_request(system_message=prompt, messages=[other, reminder, user])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert "prompt" in result.system_message.content
|
||||
assert "custom system block" in result.system_message.content
|
||||
assert "day2" in result.system_message.content
|
||||
|
||||
def test_single_reminder_not_deduplicated(self):
|
||||
"""A single reminder SystemMessage is kept — no dedup needed."""
|
||||
prompt = SystemMessage(content="prompt", id="sys-prompt")
|
||||
reminder = SystemMessage(
|
||||
content="<system-reminder>date</system-reminder>",
|
||||
id="msg-1",
|
||||
additional_kwargs={"dynamic_context_reminder": True},
|
||||
)
|
||||
user = HumanMessage(content="hi")
|
||||
request = _make_request(system_message=prompt, messages=[reminder, user])
|
||||
|
||||
result = _coalesce_request(request)
|
||||
assert result is not None
|
||||
assert "prompt" in result.system_message.content
|
||||
assert "date" in result.system_message.content
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# _flatten_content (pure helper)
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestFlattenContent:
|
||||
def test_string_content_returns_same_string(self):
|
||||
assert _flatten_content("hello") == "hello"
|
||||
|
||||
def test_list_of_strings(self):
|
||||
assert _flatten_content(["line1", "line2"]) == "line1\nline2"
|
||||
|
||||
def test_list_of_text_dicts(self):
|
||||
content = [{"type": "text", "text": "paragraph1"}, {"type": "text", "text": "paragraph2"}]
|
||||
assert _flatten_content(content) == "paragraph1\nparagraph2"
|
||||
|
||||
def test_mixed_list(self):
|
||||
content = ["plain", {"type": "text", "text": "dict"}, 42]
|
||||
assert _flatten_content(content) == "plain\ndict\n42"
|
||||
|
||||
def test_non_string_non_list(self):
|
||||
assert _flatten_content(42) == "42"
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# SystemMessageCoalescingMiddleware.wrap_model_call
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestWrapModelCall:
|
||||
def test_passthrough_when_no_in_msg_systems(self):
|
||||
"""No SystemMessages in messages → same request object passed through."""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="prompt"),
|
||||
messages=[HumanMessage(content="hi")],
|
||||
)
|
||||
captured, handler = _capture_handler()
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
assert captured[0] is request
|
||||
|
||||
def test_override_called_when_in_msg_systems_present(self):
|
||||
"""system_message + in-msg SystemMessage → override with coalesced request."""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="prompt"),
|
||||
messages=[SystemMessage(content="reminder"), HumanMessage(content="hi")],
|
||||
)
|
||||
captured, handler = _capture_handler()
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
sent = captured[0]
|
||||
assert sent is not request # overridden
|
||||
assert sent.system_message is not None
|
||||
assert "prompt" in sent.system_message.content
|
||||
assert "reminder" in sent.system_message.content
|
||||
assert not any(isinstance(m, SystemMessage) for m in sent.messages)
|
||||
|
||||
def test_returns_handler_result(self):
|
||||
"""wrap_model_call returns whatever the handler returns."""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="prompt"),
|
||||
messages=[SystemMessage(content="reminder")],
|
||||
)
|
||||
handler = MagicMock(return_value="llm-response")
|
||||
|
||||
result = mw.wrap_model_call(request, handler)
|
||||
assert result == "llm-response"
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# SystemMessageCoalescingMiddleware.awrap_model_call
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestAwrapModelCall:
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_passthrough_no_in_msg_systems(self):
|
||||
"""Async path: no in-msg systems → passthrough."""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="prompt"),
|
||||
messages=[HumanMessage(content="hi")],
|
||||
)
|
||||
captured = []
|
||||
|
||||
async def handler(req):
|
||||
captured.append(req)
|
||||
return "async-response"
|
||||
|
||||
result = await mw.awrap_model_call(request, handler)
|
||||
assert result == "async-response"
|
||||
assert captured[0] is request
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_coalesces_in_msg_systems(self):
|
||||
"""Async path: system_message + in-msg system → coalesced."""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="prompt"),
|
||||
messages=[SystemMessage(content="reminder"), HumanMessage(content="hi")],
|
||||
)
|
||||
captured = []
|
||||
|
||||
async def handler(req):
|
||||
captured.append(req)
|
||||
return "ok"
|
||||
|
||||
result = await mw.awrap_model_call(request, handler)
|
||||
assert result == "ok"
|
||||
sent = captured[0]
|
||||
assert sent is not request
|
||||
assert "prompt" in sent.system_message.content
|
||||
assert "reminder" in sent.system_message.content
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Realistic scenario: DynamicContextMiddleware ID-swap + create_agent prompt
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class TestRealisticScenario:
|
||||
def test_first_turn_single_system_message_in_final_payload(self):
|
||||
"""Simulate the exact #3707 trigger with the real request shape.
|
||||
|
||||
DynamicContextMiddleware injects the reminder into state["messages"]
|
||||
(which becomes request.messages). create_agent holds system_prompt in
|
||||
request.system_message. Without coalescing, the handler produces
|
||||
[system_prompt, reminder, __memory, __user] → 2 SystemMessages → 400.
|
||||
|
||||
With coalescing, the final payload [system_message, *messages] has
|
||||
exactly 1 SystemMessage.
|
||||
"""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
# Real shape: system_prompt in system_message field, reminder in messages
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="You are DeerFlow 2.0, an AI assistant...", id="sys-prompt"),
|
||||
messages=[
|
||||
SystemMessage(
|
||||
content="<system-reminder>\n<current_date>2026-06-22, Monday</current_date>\n</system-reminder>",
|
||||
id="msg-1",
|
||||
additional_kwargs={"hide_from_ui": True, "dynamic_context_reminder": True},
|
||||
),
|
||||
HumanMessage(content="<memory>User prefers Python.</memory>", id="msg-1__memory"),
|
||||
HumanMessage(content="What is the capital of France?", id="msg-1__user"),
|
||||
],
|
||||
)
|
||||
captured, handler = _capture_handler()
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
# Simulate what the LLM receives: [system_message, *messages]
|
||||
sent = captured[0]
|
||||
final = _final_payload(sent)
|
||||
system_count = sum(1 for m in final if isinstance(m, SystemMessage))
|
||||
assert system_count == 1 # key assertion: only 1 SystemMessage
|
||||
assert "DeerFlow 2.0" in final[0].content
|
||||
assert "<current_date>2026-06-22, Monday</current_date>" in final[0].content
|
||||
# User-owned memory stays as HumanMessage (OWASP LLM01 preserved)
|
||||
assert any(isinstance(m, HumanMessage) and "User prefers Python." in m.content for m in final)
|
||||
assert any(isinstance(m, HumanMessage) and m.content == "What is the capital of France?" for m in final)
|
||||
|
||||
def test_midnight_crossing_single_system_message_in_final_payload(self):
|
||||
"""Midnight crossing: 3 SystemMessages total → coalesced to 1.
|
||||
|
||||
DynamicContextMiddleware injects date reminders marked with
|
||||
``dynamic_context_reminder=True``. On midnight crossings a second
|
||||
reminder (day2) appears after Human/AI turns. During coalescing the
|
||||
intervening turns that originally separated the two reminders are
|
||||
stripped from the merged block, so two contradictory <current_date>
|
||||
blocks would appear adjacent. The middleware keeps only the last
|
||||
reminder (day2) and drops earlier ones (day1) so the model sees a
|
||||
single unambiguous current date.
|
||||
"""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="system prompt", id="sys-prompt"),
|
||||
messages=[
|
||||
SystemMessage(
|
||||
content="<system-reminder>day1</system-reminder>",
|
||||
id="msg-1",
|
||||
additional_kwargs={"hide_from_ui": True, "dynamic_context_reminder": True},
|
||||
),
|
||||
HumanMessage(content="<memory>...</memory>", id="msg-1__memory"),
|
||||
HumanMessage(content="first question", id="msg-1__user"),
|
||||
AIMessage(content="first answer", id="ai-1"),
|
||||
SystemMessage(
|
||||
content="<system-reminder>day2</system-reminder>",
|
||||
id="msg-2",
|
||||
additional_kwargs={"hide_from_ui": True, "dynamic_context_reminder": True},
|
||||
),
|
||||
HumanMessage(content="second question", id="msg-2__user"),
|
||||
],
|
||||
)
|
||||
captured, handler = _capture_handler()
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
sent = captured[0]
|
||||
final = _final_payload(sent)
|
||||
system_count = sum(1 for m in final if isinstance(m, SystemMessage))
|
||||
assert system_count == 1
|
||||
merged = final[0]
|
||||
assert "system prompt" in merged.content
|
||||
# Only the latest date survives; the stale day1 reminder is dropped.
|
||||
assert "day2" in merged.content
|
||||
assert "day1" not in merged.content
|
||||
# Non-system messages in order
|
||||
non_system = [m for m in final if not isinstance(m, SystemMessage)]
|
||||
assert [m.content for m in non_system] == [
|
||||
"<memory>...</memory>",
|
||||
"first question",
|
||||
"first answer",
|
||||
"second question",
|
||||
]
|
||||
|
||||
def test_no_system_message_field_but_in_msg_systems(self):
|
||||
"""Edge case: system_message is None but messages has a SystemMessage.
|
||||
|
||||
The coalesced SystemMessage moves to the system_message field so the
|
||||
handler still prepends it as a leading system block.
|
||||
"""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
request = _make_request(
|
||||
system_message=None,
|
||||
messages=[
|
||||
SystemMessage(content="reminder", id="msg-1"),
|
||||
HumanMessage(content="hi"),
|
||||
],
|
||||
)
|
||||
captured, handler = _capture_handler()
|
||||
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
sent = captured[0]
|
||||
final = _final_payload(sent)
|
||||
system_count = sum(1 for m in final if isinstance(m, SystemMessage))
|
||||
assert system_count == 1
|
||||
assert "reminder" in final[0].content
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# Strict-backend stub: end-to-end test against vLLM/SGLang/Qwen rejection
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
class StrictBackendError(Exception):
|
||||
"""Simulates the 400 Bad Request from strict OpenAI-compatible backends.
|
||||
|
||||
Qwen, SGLang, and vLLM reject requests that contain more than one
|
||||
SystemMessage or where the sole SystemMessage is not at position 0.
|
||||
"""
|
||||
|
||||
|
||||
def _strict_backend_handler(request):
|
||||
"""Simulate a strict OpenAI-compatible backend (Qwen / SGLang / vLLM).
|
||||
|
||||
These backends accept exactly 0 or 1 SystemMessage, and if present it
|
||||
must be at position 0. They reject with "System message must be at the
|
||||
beginning" or "Received multiple system messages" otherwise.
|
||||
|
||||
The handler flattens the request into ``[system_message, *messages]``
|
||||
(matching ``create_agent``'s ``_execute_model_sync``) then validates.
|
||||
"""
|
||||
final = _final_payload(request)
|
||||
system_count = sum(1 for m in final if isinstance(m, SystemMessage))
|
||||
if system_count > 1:
|
||||
raise StrictBackendError("Received multiple system messages")
|
||||
if system_count == 1 and not isinstance(final[0], SystemMessage):
|
||||
raise StrictBackendError("System message must be at the beginning")
|
||||
return "ok"
|
||||
|
||||
|
||||
async def _async_strict_backend_handler(request):
|
||||
"""Async variant of ``_strict_backend_handler`` for awrap_model_call tests."""
|
||||
return _strict_backend_handler(request)
|
||||
|
||||
|
||||
class TestStrictBackendStub:
|
||||
"""End-to-end test against a strict-backend stub.
|
||||
|
||||
Builds the request the way create_agent does (prompt in system_message,
|
||||
not messages) and asserts that:
|
||||
- Without middleware, the strict stub raises StrictBackendError (#3707 bug)
|
||||
- With middleware, the strict stub accepts the request (fix confirmed)
|
||||
"""
|
||||
|
||||
def test_first_turn_without_middleware_rejects(self):
|
||||
"""Reproduce #3707: raw request → handler flattens to 2 SystemMessages → stub rejects."""
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="You are DeerFlow.", id="sys-prompt"),
|
||||
messages=[
|
||||
SystemMessage(content="<system-reminder>date</system-reminder>", id="msg-1"),
|
||||
HumanMessage(content="Hello", id="msg-1__user"),
|
||||
],
|
||||
)
|
||||
# Final payload: [sys-prompt, reminder, __user] → 2 SystemMessages
|
||||
# → strict backend rejects: multiple system messages
|
||||
with pytest.raises(StrictBackendError):
|
||||
_strict_backend_handler(request)
|
||||
|
||||
def test_first_turn_with_middleware_accepts(self):
|
||||
"""Fix confirmed: middleware coalesces → 1 SystemMessage → stub accepts."""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="You are DeerFlow.", id="sys-prompt"),
|
||||
messages=[
|
||||
SystemMessage(content="<system-reminder>date</system-reminder>", id="msg-1"),
|
||||
HumanMessage(content="Hello", id="msg-1__user"),
|
||||
],
|
||||
)
|
||||
result = mw.wrap_model_call(request, _strict_backend_handler)
|
||||
assert result == "ok"
|
||||
|
||||
def test_midnight_crossing_without_middleware_rejects(self):
|
||||
"""3 SystemMessages without middleware → stub rejects."""
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="prompt", id="sys-prompt"),
|
||||
messages=[
|
||||
SystemMessage(
|
||||
content="<system-reminder>day1</system-reminder>",
|
||||
id="msg-1",
|
||||
additional_kwargs={"dynamic_context_reminder": True},
|
||||
),
|
||||
HumanMessage(content="q1", id="msg-1__user"),
|
||||
AIMessage(content="a1", id="ai-1"),
|
||||
SystemMessage(
|
||||
content="<system-reminder>day2</system-reminder>",
|
||||
id="msg-2",
|
||||
additional_kwargs={"dynamic_context_reminder": True},
|
||||
),
|
||||
HumanMessage(content="q2", id="msg-2__user"),
|
||||
],
|
||||
)
|
||||
with pytest.raises(StrictBackendError):
|
||||
_strict_backend_handler(request)
|
||||
|
||||
def test_midnight_crossing_with_middleware_accepts(self):
|
||||
"""3 SystemMessages with middleware → coalesced to 1 → stub accepts.
|
||||
|
||||
Only the latest reminder (day2) survives in the merged content;
|
||||
the stale day1 reminder is dropped.
|
||||
"""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="prompt", id="sys-prompt"),
|
||||
messages=[
|
||||
SystemMessage(
|
||||
content="<system-reminder>day1</system-reminder>",
|
||||
id="msg-1",
|
||||
additional_kwargs={"dynamic_context_reminder": True},
|
||||
),
|
||||
HumanMessage(content="q1", id="msg-1__user"),
|
||||
AIMessage(content="a1", id="ai-1"),
|
||||
SystemMessage(
|
||||
content="<system-reminder>day2</system-reminder>",
|
||||
id="msg-2",
|
||||
additional_kwargs={"dynamic_context_reminder": True},
|
||||
),
|
||||
HumanMessage(content="q2", id="msg-2__user"),
|
||||
],
|
||||
)
|
||||
result = mw.wrap_model_call(request, _strict_backend_handler)
|
||||
assert result == "ok"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_path_with_middleware_accepts(self):
|
||||
"""Async path: middleware + strict stub → accepts."""
|
||||
mw = SystemMessageCoalescingMiddleware()
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="prompt", id="sys-prompt"),
|
||||
messages=[
|
||||
SystemMessage(content="<system-reminder>date</system-reminder>", id="msg-1"),
|
||||
HumanMessage(content="Hello", id="msg-1__user"),
|
||||
],
|
||||
)
|
||||
result = await mw.awrap_model_call(request, _async_strict_backend_handler)
|
||||
assert result == "ok"
|
||||
|
||||
def test_clean_request_no_middleware_needed(self):
|
||||
"""Only system_message field, no in-msg systems → stub accepts without middleware."""
|
||||
request = _make_request(
|
||||
system_message=SystemMessage(content="prompt", id="sys-prompt"),
|
||||
messages=[HumanMessage(content="hi")],
|
||||
)
|
||||
# No middleware needed — single SystemMessage at position 0
|
||||
result = _strict_backend_handler(request)
|
||||
assert result == "ok"
|
||||
Loading…
Add table
Add a link
Reference in a new issue