fix(subagents): inherit summarization middleware and harden step capture (#3875 Phase 3) (#4009)

Phase 3 of #3875 — subagents previously inherited none of the lead's
context-compaction, so a deep-research subagent (max_turns up to 150)
could accumulate >1M cumulative input before max_turns/timeout/token_budget
engaged, even after Phase 2's budget capped the pathological tail.

- Gate the subagent runtime chain on the SAME ``app_config.summarization.enabled``
  switch the lead reads (per maintainer guidance in #3875), via the shared
  ``create_summarization_middleware`` factory. One config covers both chains;
  no separate ``subagents.summarization`` field. No-op when summarization is
  off (factory returns None).
- ``skip_memory_flush=True`` on the subagent path: the factory otherwise
  attaches ``memory_flush_hook`` (when memory.enabled), which flushes
  pre-compaction messages into durable memory keyed by thread_id. Subagents
  share the parent's thread_id, so without skipping the hook a subagent's
  internal turns would pollute the PARENT thread's durable memory
  (#3875 Phase 3 review point).
- Harden ``capture_new_step_messages`` to tolerate history contraction:
  summarization rewrites the messages channel via
  ``RemoveMessage(id=REMOVE_ALL_MESSAGES)``, shrinking len(messages) below
  the step-capture cursor. Without a reset, every step appended after the
  compaction point was dropped until length overtook the stale cursor (#3845
  interaction, maintainer validation point (a)). Cursor now resets to the
  new tail; id/content dedup prevents re-emitting pre-compaction steps.
- Couple the DEFAULT token-budget ceiling to ``summarization.enabled``
  (#3875 Phase 3 review point): 1M when compaction is on, 2M when off
  (preserves Phase 2's deliberate headroom for summarization-off
  deep-research runs that can exceed 1M). A user-set budget (global or
  per-agent) always wins regardless of the switch. Flagged tunable.

The summarization middleware does not implement ``consume_stop_reason``, so
the Phase 2 guard-cap stop-reason channel is unaffected.

Refs: https://github.com/bytedance/deer-flow/issues/3875
This commit is contained in:
hataa 2026-07-10 11:17:35 +08:00 committed by GitHub
parent eb8eec21ce
commit 266883b3dd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 502 additions and 16 deletions

View File

@ -379,8 +379,9 @@ Proxied through nginx: `/api/langgraph/*` → Gateway LangGraph-compatible runti
**Concurrency**: `MAX_CONCURRENT_SUBAGENTS = 3` enforced by `SubagentLimitMiddleware` (truncates excess tool calls in `after_model`); default subagent timeout `subagents.timeout_seconds=1800` (30 min) and built-in `general-purpose` `max_turns=150` (raised from 100/15-min so deep-research subtasks stop hitting `GraphRecursionError` out of the box) **Concurrency**: `MAX_CONCURRENT_SUBAGENTS = 3` enforced by `SubagentLimitMiddleware` (truncates excess tool calls in `after_model`); default subagent timeout `subagents.timeout_seconds=1800` (30 min) and built-in `general-purpose` `max_turns=150` (raised from 100/15-min so deep-research subtasks stop hitting `GraphRecursionError` out of the box)
**Flow**: `task()` tool → `SubagentExecutor` → background thread → poll 5s → SSE events → result **Flow**: `task()` tool → `SubagentExecutor` → background thread → poll 5s → SSE events → result
**Events**: `task_started`, `task_running`, `task_completed`/`task_failed`/`task_timed_out` **Events**: `task_started`, `task_running`, `task_completed`/`task_failed`/`task_timed_out`
**Guardrail caps & `stop_reason` (#3875 Phase 2)**: three independent axes can end a subagent run early, and all now surface *why* through one additive field rather than a new status enum. **Turn axis**: `recursion_limit` on the subagent `run_config` equals `max_turns`, so exhausting the turn budget raises `GraphRecursionError` from `agent.astream`; `executor.py::_aexecute` catches it specifically (before the generic `except Exception`). **Token axis**: `TokenBudgetMiddleware` is attached per-agent via `build_subagent_runtime_middlewares` from `subagents.token_budget` (default 2,000,000 tokens, warn at 0.7, hard-stop at 1.0 — a backstop against a subagent that burns tokens on trivial work). It does *not* raise: at the hard-stop threshold it strips the in-flight turn's tool calls, forces `finish_reason="stop"`, and lets the run complete naturally with a final answer. **Loop axis**: `LoopDetectionMiddleware` (attached at the same point) catches repeated identical tool-call sets — or one tool *type* called many times with varying args — and its hard-stop likewise strips `tool_calls` and forces a final answer without raising, recording `loop_capped`. Each guard exposes its cap on a per-`run_id` `consume_stop_reason(run_id)` accessor; `_aexecute` collects **every** middleware with that method (duck-typed via `hasattr`, so the executor has no import coupling to the guard classes) and surfaces the first non-`None` reason — adding a future guard needs no executor change. **Surfacing**: whichever axis fired, `_aexecute` stamps a normal status plus an additive reason — `completed` + `stop_reason=token_capped|turn_capped|loop_capped` when a usable final answer (or partial recovered from the last streamed chunk via `_extract_final_result``utils/messages.py::message_content_to_text`, returning a `"No response Generated"` sentinel when no text survived) was produced; `failed` + `stop_reason=turn_capped` when nothing usable survived. `SubagentResult.stop_reason` flows through `task_tool.py::_task_result_command``format_subagent_result_message` (renders `Task Succeeded (capped: ...)` / `Task failed (capped: ...)`) and `make_subagent_additional_kwargs`, which stamps the additive `subagent_stop_reason` key alongside the normal `subagent_status`. **Why additive, not an enum**: a new status value would break v1 consumers; an optional field is ignored by older frontends and ledger readers, so the cross-language contract (`contracts/subagent_status_contract.json` v2 + `subagents/status_contract.py` + `frontend/.../subtask-result.ts`, pinned by `test_status_values_match_contract` / `test_stop_reason_values_match_contract`) stays backward-compatible. The durable delegation ledger captures `stop_reason` onto the entry and renders model-facing guidance ("hit a guardrail cap with a partial result; reuse it, retry tighter, or raise the per-agent budget (`max_turns` / `token_budget`)") so the lead reuses a capped completion knowingly instead of mistaking it for a clean one. (Phase 1 shipped this surfacing as a `MAX_TURNS_REACHED` status enum in #3949; Phase 2 replaced that enum with the additive `stop_reason` field per the agreed design — the `max_turns_reached` status value and `SubagentStatus.MAX_TURNS_REACHED` are gone.) **Guardrail caps & `stop_reason` (#3875 Phase 2)**: three independent axes can end a subagent run early, and all now surface *why* through one additive field rather than a new status enum. **Turn axis**: `recursion_limit` on the subagent `run_config` equals `max_turns`, so exhausting the turn budget raises `GraphRecursionError` from `agent.astream`; `executor.py::_aexecute` catches it specifically (before the generic `except Exception`). **Token axis**: `TokenBudgetMiddleware` is attached per-agent via `build_subagent_runtime_middlewares` from `subagents.token_budget` (default `max_tokens` **coupled to `summarization.enabled`** — 1,000,000 when subagent summarization is on, 2,000,000 when off, warn at 0.7, hard-stop at 1.0; a user-set budget always wins regardless of the switch — #3875 Phase 3; a backstop against a subagent that burns tokens on trivial work). It does *not* raise: at the hard-stop threshold it strips the in-flight turn's tool calls, forces `finish_reason="stop"`, and lets the run complete naturally with a final answer. **Loop axis**: `LoopDetectionMiddleware` (attached at the same point) catches repeated identical tool-call sets — or one tool *type* called many times with varying args — and its hard-stop likewise strips `tool_calls` and forces a final answer without raising, recording `loop_capped`. Each guard exposes its cap on a per-`run_id` `consume_stop_reason(run_id)` accessor; `_aexecute` collects **every** middleware with that method (duck-typed via `hasattr`, so the executor has no import coupling to the guard classes) and surfaces the first non-`None` reason — adding a future guard needs no executor change. **Surfacing**: whichever axis fired, `_aexecute` stamps a normal status plus an additive reason — `completed` + `stop_reason=token_capped|turn_capped|loop_capped` when a usable final answer (or partial recovered from the last streamed chunk via `_extract_final_result``utils/messages.py::message_content_to_text`, returning a `"No response Generated"` sentinel when no text survived) was produced; `failed` + `stop_reason=turn_capped` when nothing usable survived. `SubagentResult.stop_reason` flows through `task_tool.py::_task_result_command``format_subagent_result_message` (renders `Task Succeeded (capped: ...)` / `Task failed (capped: ...)`) and `make_subagent_additional_kwargs`, which stamps the additive `subagent_stop_reason` key alongside the normal `subagent_status`. **Why additive, not an enum**: a new status value would break v1 consumers; an optional field is ignored by older frontends and ledger readers, so the cross-language contract (`contracts/subagent_status_contract.json` v2 + `subagents/status_contract.py` + `frontend/.../subtask-result.ts`, pinned by `test_status_values_match_contract` / `test_stop_reason_values_match_contract`) stays backward-compatible. The durable delegation ledger captures `stop_reason` onto the entry and renders model-facing guidance ("hit a guardrail cap with a partial result; reuse it, retry tighter, or raise the per-agent budget (`max_turns` / `token_budget`)") so the lead reuses a capped completion knowingly instead of mistaking it for a clean one. (Phase 1 shipped this surfacing as a `MAX_TURNS_REACHED` status enum in #3949; Phase 2 replaced that enum with the additive `stop_reason` field per the agreed design — the `max_turns_reached` status value and `SubagentStatus.MAX_TURNS_REACHED` are gone.)
**Step capture & persistence (#3779)**: `executor.py` captures both assistant turns (`AIMessage`) **and** tool outputs (`ToolMessage`) via `subagents/step_events.py::capture_new_step_messages`, which walks the *newly-appended tail* of each `stream_mode="values"` chunk (not just `messages[-1]`) so a multi-tool-call turn — where LangGraph's `ToolNode` appends several `ToolMessage`s in one super-step — keeps every tool output instead of dropping all but the last. `runtime/runs/worker.py::_SubagentEventBuffer` additionally persists these `task_*` custom events to the `RunEventStore` as `subagent.start`/`subagent.step`/`subagent.end` (`category="subagent"`, `task_id` in `metadata`). It **batches** writes via `put_batch` (flushing on a terminal `subagent.end`, at `FLUSH_THRESHOLD` events, and in the worker's `finally`) rather than one `put()` per step, since `put()` is a documented low-frequency path (per-thread advisory lock per call) and a deep subagent (`max_turns=150`) emits hundreds of steps on the hot stream loop. `build_subagent_step` caps both the per-step `text` and each tool call's serialized `args` at `SUBAGENT_STEP_MAX_CHARS` (flagged `truncated` / `args_truncated`) so a large `write_file`/`bash` payload can't produce an unbounded row. The dedicated category keeps them out of `list_messages` (the thread feed) while `list_events` returns them for the frontend's fetch-on-expand backfill. `list_events` accepts `task_id` (filters on `metadata["task_id"]` — SQL-side in `DbRunEventStore` via `event_metadata["task_id"].as_string()`, in-memory in the JSONL/memory stores) plus an `after_seq` forward cursor, so the card pages through one subagent's steps without the run-wide `limit` truncating the tail (no schema migration: the filter rides the existing run-scoped index). `step_events.py` is a pure, unit-tested layer (`build_subagent_step` / `subagent_run_event`). **Context compaction (#3875 Phase 3)**: subagents inherit `DeerFlowSummarizationMiddleware` via `build_subagent_runtime_middlewares`, gated on the **same** `summarization.enabled` switch the lead reads (one config covers both chains; trigger/keep/model/prompt come from the shared `summarization` config so they cannot drift). The factory is called with `skip_memory_flush=True` on the subagent path: the lead's `memory_flush_hook` (attached when `memory.enabled`) flushes pre-compaction messages into durable memory keyed by `thread_id`, and subagents share the parent's `thread_id`, so without skipping the hook a subagent's internal turns would pollute the **parent** thread's durable memory. Placement differs from the lead chain (lead appends summarization *before* the guard trio; subagent appends it *after*) — benign because the middleware implements only `before_model` (compaction) with no `after_model`/`consume_stop_reason`, so it cannot disturb the Phase 2 guard-cap stop-reason channel. Compaction rewrites the messages channel via `RemoveMessage(id=REMOVE_ALL_MESSAGES)`, which shrinks `len(messages)` below the step-capture cursor mid-run; `capture_new_step_messages` (see Step capture below) resets the cursor to the new tail on contraction so steps appended after the compaction point are not silently dropped.
**Step capture & persistence (#3779)**: `executor.py` captures both assistant turns (`AIMessage`) **and** tool outputs (`ToolMessage`) via `subagents/step_events.py::capture_new_step_messages`, which walks the *newly-appended tail* of each `stream_mode="values"` chunk (not just `messages[-1]`) so a multi-tool-call turn — where LangGraph's `ToolNode` appends several `ToolMessage`s in one super-step — keeps every tool output instead of dropping all but the last. `runtime/runs/worker.py::_SubagentEventBuffer` additionally persists these `task_*` custom events to the `RunEventStore` as `subagent.start`/`subagent.step`/`subagent.end` (`category="subagent"`, `task_id` in `metadata`). It **batches** writes via `put_batch` (flushing on a terminal `subagent.end`, at `FLUSH_THRESHOLD` events, and in the worker's `finally`) rather than one `put()` per step, since `put()` is a documented low-frequency path (per-thread advisory lock per call) and a deep subagent (`max_turns=150`) emits hundreds of steps on the hot stream loop. `build_subagent_step` caps both the per-step `text` and each tool call's serialized `args` at `SUBAGENT_STEP_MAX_CHARS` (flagged `truncated` / `args_truncated`) so a large `write_file`/`bash` payload can't produce an unbounded row. The dedicated category keeps them out of `list_messages` (the thread feed) while `list_events` returns them for the frontend's fetch-on-expand backfill. `list_events` accepts `task_id` (filters on `metadata["task_id"]` — SQL-side in `DbRunEventStore` via `event_metadata["task_id"].as_string()`, in-memory in the JSONL/memory stores) plus an `after_seq` forward cursor, so the card pages through one subagent's steps without the run-wide `limit` truncating the tail (no schema migration: the filter rides the existing run-scoped index). `step_events.py` is a pure, unit-tested layer (`build_subagent_step` / `subagent_run_event`). **History contraction (#3875 Phase 3)**: `capture_new_step_messages` assumes append-only growth, but `DeerFlowSummarizationMiddleware` rewrites the messages channel via `RemoveMessage(id=REMOVE_ALL_MESSAGES)`, shrinking `len(messages)` below the cursor mid-run. On contraction (`total < processed_count`) the cursor resets to the new tail; `capture_step_message`'s id/content dedup prevents re-emitting pre-compaction steps, so steps appended after the compaction point are still captured instead of being dropped until `total` overtakes the stale cursor.
**Deferred MCP tools** (if `tool_search.enabled`): `SubagentExecutor._build_initial_state` assembles deferral after policy filtering via the shared `assemble_deferred_tools` (fail-closed), appends the `tool_search` tool, injects the `<available-deferred-tools>` section into the subagent's `SystemMessage`, and threads the setup to `_create_agent`, which attaches `McpRoutingMiddleware` (when PR1 routing metadata matches deferred tools) before `DeferredToolFilterMiddleware` through `build_subagent_runtime_middlewares(...)`. Subagents thus withhold full MCP schemas until promotion, same as the lead agent; each task run gets a fresh `ThreadState` so promotion is isolated per run **Deferred MCP tools** (if `tool_search.enabled`): `SubagentExecutor._build_initial_state` assembles deferral after policy filtering via the shared `assemble_deferred_tools` (fail-closed), appends the `tool_search` tool, injects the `<available-deferred-tools>` section into the subagent's `SystemMessage`, and threads the setup to `_create_agent`, which attaches `McpRoutingMiddleware` (when PR1 routing metadata matches deferred tools) before `DeferredToolFilterMiddleware` through `build_subagent_runtime_middlewares(...)`. Subagents thus withhold full MCP schemas until promotion, same as the lead agent; each task run gets a fresh `ThreadState` so promotion is isolated per run
**Checkpointer isolation**: Subagent graphs are compiled with `checkpointer=False` to avoid inheriting the parent run's checkpointer, since subagents are one-shot and never resume. **Checkpointer isolation**: Subagent graphs are compiled with `checkpointer=False` to avoid inheriting the parent run's checkpointer, since subagents are one-shot and never resume.

View File

@ -437,15 +437,22 @@ def create_summarization_middleware(
*, *,
app_config: Any | None = None, app_config: Any | None = None,
keep: tuple[str, int | float] | None = None, keep: tuple[str, int | float] | None = None,
skip_memory_flush: bool = False,
) -> DeerFlowSummarizationMiddleware | None: ) -> DeerFlowSummarizationMiddleware | None:
"""Create the configured summarization middleware. """Create the configured summarization middleware.
Both the lead-agent automatic path and the manual context-compaction path Both the lead-agent automatic path and the manual context-compaction path
use this factory so model resolution, hooks, prompt config, and retention use this factory so model resolution, hooks, prompt config, and retention
defaults cannot drift. defaults cannot drift.
"""
from deerflow.agents.memory.summarization_hook import memory_flush_hook
``skip_memory_flush`` omits the ``memory_flush_hook`` that otherwise
flushes pre-compaction messages into the durable memory queue. The lead
chain keeps it (research should persist); the subagent chain sets it so a
subagent's INTERNAL turns (the "Task" human message + intermediate AI/tool
turns) are not written into the PARENT thread's durable memory — the hook
is keyed by ``thread_id`` and subagents share the parent's ``thread_id``
(#3875 Phase 3 review).
"""
resolved_app_config = app_config or get_app_config() resolved_app_config = app_config or get_app_config()
config = resolved_app_config.summarization config = resolved_app_config.summarization
@ -485,7 +492,9 @@ def create_summarization_middleware(
kwargs["summary_prompt"] = config.summary_prompt kwargs["summary_prompt"] = config.summary_prompt
hooks: list[BeforeSummarizationHook] = [] hooks: list[BeforeSummarizationHook] = []
if resolved_app_config.memory.enabled: if resolved_app_config.memory.enabled and not skip_memory_flush:
from deerflow.agents.memory.summarization_hook import memory_flush_hook
hooks.append(memory_flush_hook) hooks.append(memory_flush_hook)
return DeerFlowSummarizationMiddleware( return DeerFlowSummarizationMiddleware(

View File

@ -353,7 +353,17 @@ def build_subagent_runtime_middlewares(
# builds a fresh middleware instance (see ``executor._create_agent``), so # builds a fresh middleware instance (see ``executor._create_agent``), so
# parallel subagents cannot cross-contaminate even though they share the # parallel subagents cannot cross-contaminate even though they share the
# parent thread_id/run_id in context. # parent thread_id/run_id in context.
token_budget_config = app_config.subagents.get_token_budget_for(agent_name) if agent_name is not None else app_config.subagents.token_budget #
# Default-ceiling coupling (#3875 Phase 3 review): the default ``max_tokens``
# is re-coupled to ``summarization.enabled`` — 1M when compaction is on, 2M
# when off. This ONLY applies to the default; a user-set budget (global or
# per-agent) always wins, so a deployment that pinned a value is never
# silently changed by flipping the summarization switch.
summarization_enabled = app_config.summarization.enabled
if agent_name is not None:
token_budget_config = app_config.subagents.get_token_budget_for(agent_name, summarization_enabled=summarization_enabled)
else:
token_budget_config = app_config.subagents.token_budget
if token_budget_config.enabled: if token_budget_config.enabled:
from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
@ -369,4 +379,45 @@ def build_subagent_runtime_middlewares(
middlewares.append(SafetyFinishReasonMiddleware.from_config(safety_config)) middlewares.append(SafetyFinishReasonMiddleware.from_config(safety_config))
# DeerFlowSummarizationMiddleware — subagents inherit none of the lead's
# context compaction today (#3875 Phase 3): a deep-research subagent
# (``max_turns`` up to 150) can accumulate >1M cumulative input before
# max_turns/timeout/token_budget engage, even though Phase 2's budget now
# caps the pathological tail. Gated on the SAME
# ``app_config.summarization.enabled`` switch the lead reads (per
# maintainer guidance in #3875) so a single config covers both chains —
# no separate ``subagents.summarization`` field. The shared factory
# returns ``None`` when summarization is disabled, so this is a pure
# no-op when the switch is off. Trigger/keep/model/prompt all come from
# the same ``summarization`` config the lead reads, so the two chains
# cannot drift.
#
# Placement differs from the lead chain: the lead appends summarization
# BEFORE the guard trio (loop/token/safety), here it is appended AFTER.
# This is benign — compaction runs in ``before_model`` regardless of
# relative position, and the guard middlewares account in ``after_model``
# — but noted because the relative order is not an exact mirror.
#
# ``skip_memory_flush=True``: the factory otherwise attaches
# ``memory_flush_hook`` (when ``memory.enabled``), which flushes
# pre-compaction messages into the durable memory queue keyed by
# ``thread_id``. Subagents share the parent's ``thread_id`` in context, so
# without skipping the hook a subagent's internal turns would be written
# into the PARENT thread's durable memory (#3875 Phase 3 review).
#
# The middleware rewrites history via ``RemoveMessage(id=REMOVE_ALL_MESSAGES)``,
# which shrinks the messages channel mid-run;
# ``capture_new_step_messages`` must tolerate that contraction (see
# ``step_events.py``) or it drops steps captured after the compaction
# point. It does not implement ``consume_stop_reason``, so it does not
# interfere with the Phase 2 guard-cap stop-reason channel.
from deerflow.agents.middlewares.summarization_middleware import create_summarization_middleware
summarization_middleware = create_summarization_middleware(
app_config=app_config,
skip_memory_flush=True,
)
if summarization_middleware is not None:
middlewares.append(summarization_middleware)
return middlewares return middlewares

View File

@ -9,18 +9,34 @@ from deerflow.config.token_budget_config import TokenBudgetConfig
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def default_subagent_token_budget() -> TokenBudgetConfig: def default_subagent_token_budget(*, summarization_enabled: bool = False) -> TokenBudgetConfig:
"""Default per-run token budget for subagents (#3875 Phase 2). """Default per-run token budget for subagents (#3875 Phase 2 → Phase 3 coupling).
Enabled by default so the pathological-token-burn backstop actually Enabled by default so the pathological-token-burn backstop actually
engages (per umbrella #3857 point 4 — backstops must engage, not just engages (per umbrella #3857 point 4 — backstops must engage, not just
exist). ``max_tokens`` is a deliberately loose ceiling: the reported 4.4M exist). ``max_tokens`` is **coupled to whether subagent summarization is
burn would have been cut roughly in half, while legitimate deep-research on** (#3875 Phase 3 review point):
runs (``max_turns=150``, no summarization yet) can genuinely accumulate
>1M cumulative input today. Tighten after Phase 3 lands subagent - ``summarization_enabled=True`` (Phase 3 compacts the running context
summarization. Flagged tunable in the PR description. before it reaches pathological size): **1M** tighter ceiling still
covers legitimate deep research while catching degenerate runs earlier.
- ``summarization_enabled=False``: **2M** the Phase 2 ceiling. Phase 2's
own docstring noted legitimate deep-research runs (``max_turns=150``,
no summarization) "can genuinely accumulate >1M cumulative input," so a
1M ceiling without compaction would prematurely cap them. Keeping 2M
here preserves that headroom; the tighter 1M only applies when the
compaction that justifies it is actually running.
The model-level ``default_factory`` (``SubagentsAppConfig.token_budget``)
cannot read ``summarization.enabled`` (a sibling top-level field), so it
falls back to the 2M no-compaction default; the builder
(``build_subagent_runtime_middlewares``) recomputes via
``get_token_budget_for(..., summarization_enabled=...)`` so the live value
reflects the actual switch. A user-set ``token_budget`` (global or
per-agent) always wins regardless of the switch. Flagged tunable.
""" """
return TokenBudgetConfig(enabled=True, max_tokens=2_000_000, warn_threshold=0.7) max_tokens = 1_000_000 if summarization_enabled else 2_000_000
return TokenBudgetConfig(enabled=True, max_tokens=max_tokens, warn_threshold=0.7)
class SubagentOverrideConfig(BaseModel): class SubagentOverrideConfig(BaseModel):
@ -114,6 +130,20 @@ class SubagentsAppConfig(BaseModel):
description="User-defined subagent types keyed by agent name", description="User-defined subagent types keyed by agent name",
) )
# True when ``token_budget`` was NOT explicitly provided by the user, i.e.
# the field fell back to its default_factory. ``get_token_budget_for`` uses
# this to decide whether the ceiling may be re-coupled to
# ``summarization.enabled`` (#3875 Phase 3): a user-set budget is always
# respected as-is. Set by ``__init__`` from ``model_fields_set`` and
# preserved across the app-config reload path (which drops a default
# ``token_budget`` before re-constructing — see
# ``load_subagents_config_from_dict``).
_token_budget_is_default: bool = True
def __init__(self, **data):
super().__init__(**data)
self._token_budget_is_default = "token_budget" not in self.model_fields_set
def get_timeout_for(self, agent_name: str) -> int: def get_timeout_for(self, agent_name: str) -> int:
"""Get the effective timeout for a specific agent. """Get the effective timeout for a specific agent.
@ -165,7 +195,12 @@ class SubagentsAppConfig(BaseModel):
return override.skills return override.skills
return None return None
def get_token_budget_for(self, agent_name: str) -> TokenBudgetConfig: def get_token_budget_for(
self,
agent_name: str,
*,
summarization_enabled: bool = False,
) -> TokenBudgetConfig:
"""Get the effective token-budget config for a specific agent. """Get the effective token-budget config for a specific agent.
Unlike ``max_turns``/``timeout_seconds`` (which keep a custom agent's Unlike ``max_turns``/``timeout_seconds`` (which keep a custom agent's
@ -173,10 +208,21 @@ class SubagentsAppConfig(BaseModel):
every subagent unless explicitly disabled so the per-agent override every subagent unless explicitly disabled so the per-agent override
wins when set, otherwise the global default applies to built-in AND wins when set, otherwise the global default applies to built-in AND
custom agents alike (#3875 Phase 2 / umbrella #3857 point 4). custom agents alike (#3875 Phase 2 / umbrella #3857 point 4).
``summarization_enabled`` couples the DEFAULT ceiling to whether
subagent summarization is on (#3875 Phase 3 review): 1M when
compaction is running, 2M otherwise. It ONLY affects the default
any explicitly configured ``token_budget`` (global or per-agent)
wins regardless, so a deployment that pinned a value is never
silently changed by flipping the summarization switch.
""" """
override = self.agents.get(agent_name) override = self.agents.get(agent_name)
if override is not None and override.token_budget is not None: if override is not None and override.token_budget is not None:
return override.token_budget return override.token_budget
# Only recompute when the caller is using the default (no explicit
# global token_budget was set). A user-set global is respected as-is.
if self._token_budget_is_default:
return default_subagent_token_budget(summarization_enabled=summarization_enabled)
return self.token_budget return self.token_budget
@ -191,6 +237,18 @@ def get_subagents_app_config() -> SubagentsAppConfig:
def load_subagents_config_from_dict(config_dict: dict) -> None: def load_subagents_config_from_dict(config_dict: dict) -> None:
"""Load subagents configuration from a dictionary.""" """Load subagents configuration from a dictionary."""
global _subagents_config global _subagents_config
# The app-config reload path (app_config.py) round-trips via
# ``config.subagents.model_dump()``, which serializes a default
# ``token_budget`` into the dict. Re-constructing from that dict would make
# ``model_fields_set`` contain ``token_budget`` and flip
# ``_token_budget_is_default`` to False — breaking the
# summarization-coupled recompute in ``get_token_budget_for`` (#3875 Phase
# 3). Drop the key when its value still equals the no-compaction default so
# the default_factory fires on reconstruction and the "user did not set
# this" signal is preserved.
tb = config_dict.get("token_budget")
if tb is not None and tb == default_subagent_token_budget(summarization_enabled=False).model_dump():
config_dict = {k: v for k, v in config_dict.items() if k != "token_budget"}
_subagents_config = SubagentsAppConfig(**config_dict) _subagents_config = SubagentsAppConfig(**config_dict)
overrides_summary = {} overrides_summary = {}

View File

@ -95,8 +95,27 @@ def capture_new_step_messages(
grow, re-examine only the trailing message so an id-less in-place replacement grow, re-examine only the trailing message so an id-less in-place replacement
(same length, new content) is still captured ``capture_step_message``'s (same length, new content) is still captured ``capture_step_message``'s
dedup makes an unchanged re-yield a no-op. Returns the new cursor. dedup makes an unchanged re-yield a no-op. Returns the new cursor.
When the history *contracted* (``total < processed_count``) which happens
when ``DeerFlowSummarizationMiddleware`` rewrites the channel via
``RemoveMessage(id=REMOVE_ALL_MESSAGES)`` (#3875 Phase 3) — reset the cursor
to the new tail and let ``capture_step_message``'s id/content dedup prevent
re-emitting steps captured before the compaction. Without this reset, every
step appended after the compaction point is dropped until ``total`` overtakes
the stale cursor.
INVARIANT: after the reset the no-growth branch only re-examines
``messages[-1]``, so a genuinely new AIMessage/ToolMessage inserted at an
index BELOW the reset cursor in a compacted list would be missed. This is
not reachable today: the summarization middleware puts the summary into a
separate ``summary_text`` state key, and the messages channel after
compaction holds only already-seen preserved tail messages compaction
never inserts a NEW capturable message below the cursor. If a future
middleware violates this invariant, the reset branch needs a full re-scan.
""" """
total = len(messages) total = len(messages)
if total < processed_count:
processed_count = total
if total > processed_count: if total > processed_count:
for message in messages[processed_count:total]: for message in messages[processed_count:total]:
capture_step_message(message, captured, seen_ids) capture_step_message(message, captured, seen_ids)

View File

@ -841,6 +841,60 @@ class TestAsyncExecutionPath:
assert [m["id"] for m in result.ai_messages] == ["ai-1", "tool-1", "tool-2", "tool-3", "ai-2"] assert [m["id"] for m in result.ai_messages] == ["ai-1", "tool-1", "tool-2", "tool-3", "ai-2"]
@pytest.mark.anyio
async def test_aexecute_step_capture_survives_history_contraction(self, classes, base_config, mock_agent, msg):
"""Regression for #3875 Phase 3: DeerFlowSummarizationMiddleware rewrites the
messages channel mid-run via ``RemoveMessage(id=REMOVE_ALL_MESSAGES)``,
so a later ``values`` snapshot hands the executor a SHORTER message list
than the cursor it was tracking. Without the contraction reset in
``capture_new_step_messages``, every step appended after the compaction
is dropped until the list length overtakes the stale cursor.
Faithful to the real middleware: compaction puts the summary into a
SEPARATE ``summary_text`` state key the messages channel after
compaction holds only the preserved recent tail (already-seen
messages), NOT a synthetic summary AIMessage. So the contraction chunk
is the already-seen tail (deduped, no new step); the real regression
coverage is that POST-compaction growth is still captured."""
SubagentExecutor = classes["SubagentExecutor"]
human = msg.human("Task")
ai1 = msg.ai("turn one", "ai-1")
tool1 = msg.tool("r1", "call_1", name="web_search", msg_id="tool-1")
ai2 = msg.ai("turn two", "ai-2") # also the preserved tail after compaction
tool2 = msg.tool("r2", "call_2", name="read_file", msg_id="tool-2")
final = msg.ai("final answer", "ai-3")
chunks = [
# Pre-compaction growth (cursor → 4).
{"messages": [human, ai1]},
{"messages": [human, ai1, tool1]},
{"messages": [human, ai1, tool1, ai2]},
# Compaction: channel rewrites to just the preserved tail (ai2) —
# length drops from 4 to 1, below the cursor. ai2 is already seen
# (deduped), so no new step is emitted. (The summary lives in
# summary_text, out of channel.)
{"messages": [ai2]},
# Post-compaction growth — the bug: tool-2/final were dropped.
{"messages": [ai2, tool2]},
{"messages": [ai2, tool2, final]},
]
mock_agent.astream = lambda *args, **kwargs: async_iterator(chunks)
executor = SubagentExecutor(config=base_config, tools=[], thread_id="test-thread")
with patch.object(executor, "_create_agent", return_value=mock_agent):
result = await executor._aexecute("Task")
# Pre-compaction steps survive (ai2 not re-emitted — deduped), and
# crucially the post-compaction tool + final answer are NOT dropped.
assert [m["id"] for m in result.ai_messages] == [
"ai-1",
"tool-1",
"ai-2",
"tool-2",
"ai-3",
]
@pytest.mark.anyio @pytest.mark.anyio
async def test_aexecute_handles_list_content(self, classes, base_config, mock_agent, msg): async def test_aexecute_handles_list_content(self, classes, base_config, mock_agent, msg):
"""Test handling of list-type content in AIMessage.""" """Test handling of list-type content in AIMessage."""

View File

@ -243,6 +243,52 @@ def test_capture_new_step_messages_is_noop_on_values_reyield():
assert len(captured) == 1 assert len(captured) == 1
def test_capture_new_step_messages_handles_history_contraction():
# Regression for #3875 Phase 3: DeerFlowSummarizationMiddleware rewrites the
# messages channel via RemoveMessage(id=REMOVE_ALL_MESSAGES), which shrinks
# len(messages) below the cursor we were tracking. Without a contraction
# reset, every step appended AFTER the compaction is dropped until total
# overtakes the stale cursor.
#
# Faithful to the real middleware: compaction puts the summary into a
# SEPARATE ``summary_text`` state key — the messages channel after
# compaction holds only the preserved recent tail (already-seen messages),
# NOT a synthetic summary AIMessage. So the contraction chunk is the
# already-seen tail, deduped by id; the real regression coverage is that
# POST-compaction growth is still captured.
captured: list[dict] = []
seen: set[str] = set()
# Pre-compaction: a normal growing turn captures 3 steps (cursor → 4).
ai1 = AIMessage(content="searching", id="ai-1")
tool1 = ToolMessage(content="r1", tool_call_id="c1", name="web_search", id="tool-1")
ai2 = AIMessage(content="done turn", id="ai-2")
before = [HumanMessage(content="do research", id="h-1"), ai1, tool1, ai2]
processed = capture_new_step_messages(before, captured, seen, 0)
assert processed == 4
assert [c["id"] for c in captured] == ["ai-1", "tool-1", "ai-2"]
# Compaction rewrites the channel to just the preserved tail (ai2) —
# length drops from 4 to 1, below the cursor. ai2 is already seen, so the
# dedup makes it a no-op; no new step is emitted. (The summary itself lives
# in summary_text and is never a capturable AIMessage — see step_events.py
# INVARIANT.)
compacted = [ai2]
processed = capture_new_step_messages(compacted, captured, seen, processed)
assert processed == 1
assert [c["id"] for c in captured] == ["ai-1", "tool-1", "ai-2"] # unchanged
# Post-compaction growth: a new turn appends after the preserved tail. This
# is the bug the fix targets — without the reset, processed_count stays at
# 4, total (3) never exceeds it, and tool-2/ai-3 are silently dropped.
tool2 = ToolMessage(content="r2", tool_call_id="c2", name="read_file", id="tool-2")
ai3 = AIMessage(content="final answer", id="ai-3")
after = [ai2, tool2, ai3]
processed = capture_new_step_messages(after, captured, seen, processed)
assert processed == 3
assert [c["id"] for c in captured] == ["ai-1", "tool-1", "ai-2", "tool-2", "ai-3"]
def test_run_event_for_task_started(): def test_run_event_for_task_started():
record = subagent_run_event({"type": "task_started", "task_id": "call_1", "description": "research X"}) record = subagent_run_event({"type": "task_started", "task_id": "call_1", "description": "research X"})

View File

@ -14,6 +14,7 @@ import pytest
from deerflow.config.subagents_config import ( from deerflow.config.subagents_config import (
SubagentOverrideConfig, SubagentOverrideConfig,
SubagentsAppConfig, SubagentsAppConfig,
default_subagent_token_budget,
get_subagents_app_config, get_subagents_app_config,
load_subagents_config_from_dict, load_subagents_config_from_dict,
) )
@ -124,6 +125,57 @@ class TestSubagentsAppConfigDefaults:
with pytest.raises(ValueError): with pytest.raises(ValueError):
SubagentsAppConfig(max_turns=-60) SubagentsAppConfig(max_turns=-60)
def test_default_token_budget_coupled_to_summarization_switch(self):
"""The token-budget backstop engages by default (#3857 point 4). Its
``max_tokens`` ceiling is coupled to whether subagent summarization is
on (#3875 Phase 3 review): 1M when compaction runs, 2M otherwise —
because Phase 2 acknowledged legitimate deep-research runs can exceed
1M without compaction, so tightening to 1M unconditionally would
prematurely cap summarization-off deployments."""
# Default (no summarization): 2M — preserves Phase 2 headroom.
budget_off = default_subagent_token_budget(summarization_enabled=False)
assert budget_off.enabled is True
assert budget_off.max_tokens == 2_000_000
assert budget_off.warn_threshold == 0.7
# Summarization on: 1M — compaction justifies the tighter ceiling.
budget_on = default_subagent_token_budget(summarization_enabled=True)
assert budget_on.max_tokens == 1_000_000
# The AppConfig model-level default cannot read summarization.enabled,
# so it falls back to the no-compaction 2M; the builder recomputes via
# get_token_budget_for(summarization_enabled=...).
config = SubagentsAppConfig()
assert config.token_budget.max_tokens == 2_000_000
assert config.token_budget.enabled is True
def test_get_token_budget_for_couples_default_to_summarization(self):
"""``get_token_budget_for`` must re-couple the DEFAULT ceiling to the
summarization switch, but a user-set budget (global or per-agent) must
always win regardless of the switch (#3875 Phase 3 review)."""
# Default global budget → re-coupled.
config = SubagentsAppConfig()
assert config.get_token_budget_for("general-purpose", summarization_enabled=True).max_tokens == 1_000_000
assert config.get_token_budget_for("general-purpose", summarization_enabled=False).max_tokens == 2_000_000
def test_get_token_budget_for_respects_explicit_global(self):
"""A user-set global ``token_budget`` is respected as-is — the
summarization coupling only affects the default."""
from deerflow.config.token_budget_config import TokenBudgetConfig
config = SubagentsAppConfig(token_budget=TokenBudgetConfig(enabled=True, max_tokens=500_000))
# Explicit global wins for an agent with no per-agent override.
assert config.get_token_budget_for("general-purpose", summarization_enabled=True).max_tokens == 500_000
assert config.get_token_budget_for("general-purpose", summarization_enabled=False).max_tokens == 500_000
def test_get_token_budget_for_respects_per_agent_override(self):
"""A per-agent ``token_budget`` override wins over both the default and
the summarization coupling."""
from deerflow.config.token_budget_config import TokenBudgetConfig
config = SubagentsAppConfig(
agents={"bash": SubagentOverrideConfig(token_budget=TokenBudgetConfig(enabled=True, max_tokens=300_000))},
)
assert config.get_token_budget_for("bash", summarization_enabled=True).max_tokens == 300_000
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# SubagentsAppConfig resolution helpers # SubagentsAppConfig resolution helpers

View File

@ -13,9 +13,10 @@ from langgraph.constants import TAG_NOSTREAM
from deerflow.agents.memory.summarization_hook import memory_flush_hook from deerflow.agents.memory.summarization_hook import memory_flush_hook
from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, DynamicContextMiddleware, is_dynamic_context_reminder from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, DynamicContextMiddleware, is_dynamic_context_reminder
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware, SummarizationEvent from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware, SummarizationEvent, create_summarization_middleware
from deerflow.agents.thread_state import ThreadState from deerflow.agents.thread_state import ThreadState
from deerflow.config.memory_config import MemoryConfig from deerflow.config.memory_config import MemoryConfig
from deerflow.config.summarization_config import SummarizationConfig
def _messages() -> list: def _messages() -> list:
@ -629,3 +630,42 @@ def test_multiple_id_swap_triplets_preserve_chronological_order() -> None:
f"{base2}__memory", f"{base2}__memory",
f"{base2}__user", f"{base2}__user",
] ]
def test_factory_attaches_memory_flush_hook_by_default(monkeypatch):
"""The lead path keeps ``memory_flush_hook`` so pre-compaction messages
persist into durable memory. Verified via the factory with memory enabled
and the default ``skip_memory_flush=False``."""
fake_model = MagicMock()
fake_model.with_config.return_value = fake_model
monkeypatch.setattr("deerflow.agents.middlewares.summarization_middleware.create_chat_model", lambda **kw: fake_model)
app_config = SimpleNamespace(
summarization=SummarizationConfig(enabled=True),
memory=MemoryConfig(enabled=True),
)
middleware = create_summarization_middleware(app_config=app_config)
assert middleware is not None
assert memory_flush_hook in middleware._before_summarization_hooks
def test_factory_skip_memory_flush_omits_hook(monkeypatch):
"""``skip_memory_flush=True`` (the subagent path) must omit
``memory_flush_hook``: subagents share the parent's ``thread_id``, so
without skipping the hook a subagent's internal turns would flush into the
PARENT thread's durable memory (#3875 Phase 3 review)."""
fake_model = MagicMock()
fake_model.with_config.return_value = fake_model
monkeypatch.setattr("deerflow.agents.middlewares.summarization_middleware.create_chat_model", lambda **kw: fake_model)
app_config = SimpleNamespace(
summarization=SummarizationConfig(enabled=True),
memory=MemoryConfig(enabled=True),
)
middleware = create_summarization_middleware(app_config=app_config, skip_memory_flush=True)
assert middleware is not None
# memory.enabled is True but the hook is skipped — the whole point.
assert memory_flush_hook not in middleware._before_summarization_hooks
assert middleware._before_summarization_hooks == []

View File

@ -628,6 +628,57 @@ def test_subagent_runtime_middlewares_place_loop_detection_before_safety_finish(
assert loop_idx < safety_idx assert loop_idx < safety_idx
def test_subagent_runtime_middlewares_attach_summarization_when_enabled(monkeypatch):
"""Subagents must inherit the lead's DeerFlowSummarizationMiddleware so a
long-running deep-research subagent compacts its context before it
accumulates pathological input (#3875 Phase 3). Gated on the SAME
``summarization.enabled`` switch the lead reads one config covers both
chains. The factory is patched at its source module so the test does not
depend on a real chat model; only the gating + wiring is asserted."""
from deerflow.agents.middlewares import summarization_middleware as sm
sentinel = object()
captured: dict[str, object] = {}
def fake_create_summarization_middleware(*, app_config=None, keep=None, skip_memory_flush=False):
captured["app_config"] = app_config
captured["keep"] = keep
captured["skip_memory_flush"] = skip_memory_flush
return sentinel
# summarization is enabled by default False; flip it on so the factory path
# is taken (the factory early-returns None when disabled).
from deerflow.config.summarization_config import SummarizationConfig
app_config = _make_app_config().model_copy(update={"summarization": SummarizationConfig(enabled=True)})
monkeypatch.setattr(sm, "create_summarization_middleware", fake_create_summarization_middleware)
_stub_runtime_middleware_imports(monkeypatch)
middlewares = build_subagent_runtime_middlewares(app_config=app_config, model_name="test-model")
# The shared factory received the same app_config the builder did (no lead
# wrapper, no config drift between the two chains).
assert captured["app_config"] is app_config
# skip_memory_flush=True so subagent-internal turns are not flushed into the
# PARENT thread's durable memory (#3875 Phase 3 review).
assert captured["skip_memory_flush"] is True
assert sentinel in middlewares
def test_subagent_runtime_middlewares_omit_summarization_when_factory_returns_none(monkeypatch):
"""When ``summarization.enabled`` is False the shared factory returns None and
the subagent chain must NOT carry a summarization middleware the default
state, since SummarizationConfig.enabled defaults to False."""
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware
app_config = _make_app_config() # summarization.enabled defaults to False
_stub_runtime_middleware_imports(monkeypatch)
middlewares = build_subagent_runtime_middlewares(app_config=app_config, model_name="test-model")
assert not any(isinstance(m, DeerFlowSummarizationMiddleware) for m in middlewares)
def test_lead_runtime_chain_finds_historical_uploads_under_lazy_init_false(tmp_path, monkeypatch): def test_lead_runtime_chain_finds_historical_uploads_under_lazy_init_false(tmp_path, monkeypatch):
"""Integration anchor for the ThreadData → Uploads ordering. """Integration anchor for the ThreadData → Uploads ordering.
@ -680,3 +731,108 @@ def test_lead_runtime_chain_finds_historical_uploads_under_lazy_init_false(tmp_p
assert "<uploaded_files>" in injected_content assert "<uploaded_files>" in injected_content
assert "prior-report.txt" in injected_content assert "prior-report.txt" in injected_content
assert "previous messages" in injected_content # historical section header assert "previous messages" in injected_content # historical section header
def test_subagent_summarization_fires_mid_run_and_produces_usable_result(monkeypatch):
"""Integration coverage for #3875 Phase 3 review gap: drive the REAL
``DeerFlowSummarizationMiddleware`` (the exact instance the subagent chain
gets via ``create_summarization_middleware(skip_memory_flush=True)``) through
a ``create_agent`` run, and assert that (a) compaction actually fires mid-run
(messages channel contracts via ``RemoveMessage``) and (b) the run still
completes with a usable final answer not just wiring.
The builder-wiring test above proves the middleware lands on the chain; this
proves the live middleware triggers and the run survives it. We bypass the
full ``build_subagent_runtime_middlewares`` chain (whose sandbox/thread-data
stubs aren't AgentMiddleware-compatible for a live run) and use the factory
directly the same instance the builder appends."""
from langchain.agents import create_agent
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from deerflow.agents.middlewares.summarization_middleware import (
DeerFlowSummarizationMiddleware,
create_summarization_middleware,
)
from deerflow.agents.thread_state import ThreadState
from deerflow.config.memory_config import MemoryConfig
from deerflow.config.summarization_config import ContextSize, SummarizationConfig
# A model that always emits a plain AIMessage — no tools, so the run is a
# single turn but the input already exceeds the trigger threshold, forcing
# before_model compaction on the first (and only) model call.
class _StaticModel(BaseChatModel):
text: str = "final answer after compaction"
@property
def _llm_type(self) -> str:
return "static"
def bind_tools(self, tools, **kwargs):
return self
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
return ChatResult(generations=[ChatGeneration(message=AIMessage(content=self.text))])
static_model = _StaticModel()
# The factory resolves its summary model via create_chat_model; point it at
# the same static model so no real provider is contacted.
monkeypatch.setattr(
"deerflow.agents.middlewares.summarization_middleware.create_chat_model",
lambda **kwargs: static_model,
)
app_config = SimpleNamespace(
summarization=SummarizationConfig(
enabled=True,
trigger=ContextSize(type="messages", value=4),
keep=ContextSize(type="messages", value=2),
),
# memory disabled + skip_memory_flush=True mirrors the subagent path:
# no memory_flush_hook is attached.
memory=MemoryConfig(enabled=False),
)
middleware = create_summarization_middleware(
app_config=app_config,
skip_memory_flush=True,
)
assert isinstance(middleware, DeerFlowSummarizationMiddleware), "the real middleware must be built"
# Subagent invariant: skip_memory_flush means no durable-memory hook.
assert not middleware._before_summarization_hooks
agent = create_agent(
model=static_model,
tools=[],
middleware=[middleware],
state_schema=ThreadState,
)
# 6 messages > trigger(4) → compaction must fire in before_model.
seed = [
HumanMessage(content="q1", id="h1"),
AIMessage(content="a1", id="a1"),
HumanMessage(content="q2", id="h2"),
AIMessage(content="a2", id="a2"),
HumanMessage(content="q3", id="h3"),
AIMessage(content="a3", id="a3"),
]
chunks = list(agent.stream({"messages": seed}, stream_mode="updates"))
# (a) Compaction fired: the middleware's before_model emitted a summary + RemoveMessage.
before_model_chunks = [c for c in chunks if "DeerFlowSummarizationMiddleware.before_model" in c]
assert before_model_chunks, "summarization before_model must fire when messages exceed the trigger"
summary_update = before_model_chunks[0]["DeerFlowSummarizationMiddleware.before_model"]
assert summary_update.get("summary_text"), "a summary must be produced"
emitted = summary_update["messages"]
assert isinstance(emitted[0], RemoveMessage), "compaction must lead with RemoveMessage"
# (b) The run completed with a usable final AIMessage despite compaction.
# The model's output surfaces under the "model" node key in updates mode.
final_messages: list = []
for chunk in chunks:
node_msg = chunk.get("model") or chunk.get("agent") or {}
final_messages = node_msg.get("messages", final_messages)
ai_finals = [m for m in final_messages if isinstance(m, AIMessage)]
assert ai_finals, "the run must produce a final AIMessage after compaction"
assert ai_finals[-1].content == "final answer after compaction"