From 67ffbfdbb5847bff7a34d9127841be38145eb854 Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Sun, 27 Sep 2026 21:04:39 +0800 Subject: [PATCH] fix(runtime): drain subagent event flush on cancellation (#5950) * fix(runtime): drain subagent event flush on cancellation * test(runtime): pin subagent event flush cancellation ownership * docs(subagents): pin event flush cancellation ownership * docs: trim subagent guidance to fit inherited budget --------- Co-authored-by: Willem Jiang --- .../harness/deerflow/runtime/runs/worker.py | 35 ++++++++- .../harness/deerflow/subagents/AGENTS.md | 2 +- ...orker_subagent_event_flush_cancellation.py | 75 +++++++++++++++++++ 3 files changed, 107 insertions(+), 5 deletions(-) create mode 100644 backend/tests/test_worker_subagent_event_flush_cancellation.py diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index f1b1d593d8..500a7a4fe0 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -773,13 +773,40 @@ class _SubagentEventBuffer: return batch = self._pending self._pending = [] - try: - await self._event_store.put_batch(batch) - except Exception: + write_task = asyncio.create_task(self._event_store.put_batch(batch)) + cancellation: asyncio.CancelledError | None = None + failure: Exception | None = None + while True: + try: + await asyncio.shield(write_task) + break + except asyncio.CancelledError as exc: + host_task = asyncio.current_task() + if host_task is None or not host_task.cancelling(): + # The store task itself was cancelled. It did not confirm + # durability, so retain the batch for a later flush. + self._pending = batch + self._pending + raise + if cancellation is None: + cancellation = exc + while host_task.cancelling(): + host_task.uncancel() + except Exception as exc: + failure = exc + break + + if failure is not None: # Re-buffer the failed batch (ahead of any events queued since) so a # transient store error does not silently drop subagent step events. self._pending = batch + self._pending - logger.warning("Run %s: failed to persist %d subagent step event(s)", self._run_id, len(batch), exc_info=True) + logger.warning( + "Run %s: failed to persist %d subagent step event(s)", + self._run_id, + len(batch), + exc_info=(type(failure), failure, failure.__traceback__), + ) + if cancellation is not None: + raise cancellation def _bind_trace_id(config: dict[str, Any], runtime_ctx: dict[str, Any]) -> str: diff --git a/backend/packages/harness/deerflow/subagents/AGENTS.md b/backend/packages/harness/deerflow/subagents/AGENTS.md index 7d193a243e..7d862ff5a6 100644 --- a/backend/packages/harness/deerflow/subagents/AGENTS.md +++ b/backend/packages/harness/deerflow/subagents/AGENTS.md @@ -43,7 +43,7 @@ does not change retry policy. **Handled LLM failures**: `LLMErrorHandlingMiddleware` deliberately converts provider/model exceptions into an `AIMessage` so the graph can end cleanly, stamping `additional_kwargs.deerflow_error_fallback=true` plus error metadata. Clean graph termination does not imply subagent success: `SubagentExecutor` inspects the last assistant message at terminalization and maps a marked fallback to `SubagentStatus.FAILED`, which then emits `task_failed` and the existing structured `subagent_error`. Only the marker is authoritative — error-looking assistant prose without it remains a normal completed result, so neither the executor nor frontend parses display text as a status protocol. **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**: `max_turns` counts turns, `recursion_limit` counts super-steps — one per graph node, and `create_agent` compiles a node per middleware hook — so `turn_budget.py::resolve_recursion_limit` scales the turn budget by the assembled chain's depth (`test_subagent_turn_budget.py` pins it against the real compiled graph); a verbatim hand-off bought `general-purpose` ~18 of its 150 turns. Exhausting it 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.) **Context compaction (#3875 Phase 3, #4039)**: 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 subagent builder attaches `DurableContextMiddleware` immediately before summarization, using the same skills path/read-tool settings as the lead chain. Compaction stores the generated summary in `ThreadState.summary_text` rather than as a `messages` item; the durable-context wrapper therefore projects it into the next model request as guarded hidden human data. This is required when a message-count keep policy preserves only an assistant tool-call plus its tool results: without the injected summary the next request begins with assistant/tool history and strict OpenAI-compatible providers can reject it. Because `DurableContextMiddleware` inserts a second `SystemMessage(authority_contract)` after the subagent's leading system prompt, the builder also appends `SystemMessageCoalescingMiddleware` innermost (mirroring the lead chain, appended after the optional summarization middleware so it is unconditionally last) to merge every `SystemMessage` into one leading `system_message` — otherwise the durable fix would trade #4039's assistant-first HTTP 400 for a duplicate-system 400 on the same strict backends (#4040). 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. `subagent_run_event` rejects malformed chunks that lack a non-empty `task_id`; running chunks additionally require a non-negative integer `message_index` and a message object, so persisted records always satisfy the required lifecycle envelope. `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. +**Step capture & persistence (#3779)**: `executor.py` uses `step_events.py::capture_new_step_messages` to capture every new `AIMessage` and `ToolMessage` in each `stream_mode="values"` tail, including all outputs of a multi-call `ToolNode` step. `runtime/runs/worker.py::_SubagentEventBuffer` persists `task_*` events as `subagent.start`/`subagent.step`/`subagent.end` in `RunEventStore` (`category="subagent"`, `task_id` in `metadata`). Use `put_batch`, flushing at terminal `subagent.end`, `FLUSH_THRESHOLD`, and worker `finally`; per-step `put()` incurs a per-thread advisory lock on the hot stream loop. An admitted batch write drains across repeated host cancellation; if the drained write fails, its batch is re-buffered before the original cancellation propagates. `subagent_run_event` rejects malformed chunks that lack a non-empty `task_id`; running chunks additionally require a non-negative integer `message_index` and a message object, so persisted records always satisfy the required lifecycle envelope. `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` applies the subagent name allow/deny list and assembly-time authorization before calling the shared `assemble_deferred_tools`, appends the `tool_search` tool, injects the `` 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(...)`. Runtime skill policy is intentionally later and dynamic: `tool_search` may disclose/promote catalog metadata, but `SkillToolPolicyMiddleware` still removes or blocks any promoted business tool omitted by the active skill. 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. **Checkpoint lineage / stream isolation**: `_aexecute` deliberately omits checkpoint-coordinate keys (`thread_id`, `checkpoint_ns`, `checkpoint_id`, `checkpoint_map`) from the child `RunnableConfig`. LangGraph must inherit those coordinates from the copied parent ContextVar so the delegated graph retains a non-root subgraph namespace; explicitly re-supplying even the same parent `thread_id` starts a new root lineage on LangGraph 1.2.6+ and can route child AI/tool frames into the parent `messages` stream. DeerFlow business components still receive the parent `thread_id` through `runtime.context`, which is the preferred lookup path for sandbox, middleware, and attribution code. Regression coverage in `tests/test_subagent_executor.py::TestSubagentCheckpointLineage` keeps the invocation-contract assertion active on every supported version and version-gates the production-shaped parent-stream test to LangGraph 1.2.6+, where the leak exists. diff --git a/backend/tests/test_worker_subagent_event_flush_cancellation.py b/backend/tests/test_worker_subagent_event_flush_cancellation.py new file mode 100644 index 0000000000..fdb442af4f --- /dev/null +++ b/backend/tests/test_worker_subagent_event_flush_cancellation.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from deerflow.runtime.runs.worker import _SubagentEventBuffer + + +class _BlockingStore: + def __init__(self, *, fail: bool = False) -> None: + self.entered = asyncio.Event() + self.release = asyncio.Event() + self.fail = fail + self.batches: list[list[dict]] = [] + + async def put_batch(self, events): + self.entered.set() + await self.release.wait() + if self.fail: + raise RuntimeError("event store write failed") + self.batches.append([dict(event) for event in events]) + return list(events) + + +def _buffer(store: _BlockingStore) -> _SubagentEventBuffer: + buffer = _SubagentEventBuffer(store, "thread-1", "run-1") + buffer._pending.append( + { + "thread_id": "thread-1", + "run_id": "run-1", + "event_type": "subagent.step", + "category": "subagent", + "content": {"task_id": "task-1"}, + } + ) + return buffer + + +@pytest.mark.asyncio +async def test_repeated_cancellation_drains_inflight_subagent_event_flush() -> None: + store = _BlockingStore() + buffer = _buffer(store) + flush_task = asyncio.create_task(buffer.flush()) + + await store.entered.wait() + flush_task.cancel() + await asyncio.sleep(0) + flush_task.cancel() + await asyncio.sleep(0) + + assert not flush_task.done() + store.release.set() + with pytest.raises(asyncio.CancelledError): + await flush_task + + assert len(store.batches) == 1 + assert buffer._pending == [] + + +@pytest.mark.asyncio +async def test_cancelled_failed_subagent_event_flush_rebuffers_before_propagating() -> None: + store = _BlockingStore(fail=True) + buffer = _buffer(store) + original = list(buffer._pending) + flush_task = asyncio.create_task(buffer.flush()) + + await store.entered.wait() + flush_task.cancel() + await asyncio.sleep(0) + store.release.set() + + with pytest.raises(asyncio.CancelledError): + await flush_task + assert buffer._pending == original