fix(subagents): close stream before releasing resources (#5221)

Co-authored-by: CorgiBoyG <CorgiBoyG@users.noreply.github.com>
This commit is contained in:
RongJie G 2026-09-18 14:23:06 +08:00 committed by GitHub
parent cc27730348
commit 94110e5dce
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 587 additions and 53 deletions

View File

@ -0,0 +1,14 @@
from __future__ import annotations
import inspect
from typing import Any
async def close_agent_stream(stream: Any) -> None:
"""Close an agent stream when its runtime exposes asynchronous cleanup."""
close = getattr(stream, "aclose", None)
if close is None:
return
result = close()
if inspect.isawaitable(result):
await result

View File

@ -80,6 +80,7 @@ from deerflow.runtime.goal import (
write_thread_goal,
)
from deerflow.runtime.keyed_lock import AsyncKeyedLockTable
from deerflow.runtime.runs.stream_cleanup import close_agent_stream
from deerflow.runtime.serialization import serialize
from deerflow.runtime.stream_bridge import StreamBridge
from deerflow.runtime.stream_modes import normalize_stream_modes, to_langgraph_stream_modes
@ -169,16 +170,6 @@ def _schedule_terminal_cycle_collection() -> None:
loop.call_later(delay, _start_collection, context=Context())
async def _close_agent_stream(stream: Any) -> None:
"""Close a LangGraph stream deterministically after completion or early exit."""
close = getattr(stream, "aclose", None)
if close is None:
return
result = close()
if inspect.isawaitable(result):
await result
def _remove_callback(config: dict[str, Any], handler: Any) -> None:
callbacks = config.get("callbacks")
if isinstance(callbacks, list):
@ -1281,7 +1272,7 @@ async def run_agent(
finally:
close_error = sys.exception()
try:
await _close_agent_stream(stream)
await close_agent_stream(stream)
except Exception:
abort_requested = broke_on_abort or record.abort_event.is_set()
if close_error is None and not abort_requested:
@ -1330,7 +1321,7 @@ async def run_agent(
finally:
close_error = sys.exception()
try:
await _close_agent_stream(stream)
await close_agent_stream(stream)
except Exception:
abort_requested = broke_on_abort or record.abort_event.is_set()
if close_error is None and not abort_requested:

View File

@ -21,7 +21,7 @@ executions are not checked, and acceptance never changes automatic retry policy.
**Upload-state boundary**: Ordinary `task` delegation snapshots a valid parent `ThreadState.uploaded_files` list at dispatch, deep-copies it across the isolated-loop boundary, seeds it into the child's fresh state, and only then makes `list_uploaded_files` eligible for normal tool-policy filtering. An explicit empty list is valid and must be preserved because it means every upload in the thread is historical for this run. Missing or malformed state fails closed with the tool disabled. Durable `batch_task` execution intentionally keeps the tool disabled: delayed and recovered items have no valid parent-run upload boundary, and supporting that case requires a separate persisted-state contract.
**Date context (#4781)**: Every built-in subagent execution registers `SubagentDateContextMiddleware` immediately before `SystemMessageCoalescingMiddleware`. Its one-time `before_agent` hook adds a hidden framework-owned `SystemMessage` containing only `<current_date>` before the first model call; it does not read `AppConfig.memory`, call the memory manager, rewrite the task `HumanMessage`, or inherit the lead agent's frozen-conversation/midnight lifecycle. The coalescer merges that reminder with the subagent's static prompt so strict providers still receive exactly one leading `SystemMessage`. The lead-only `DynamicContextMiddleware` registration and its date, optional-memory, and midnight-update behavior remain unchanged.
**Execution**: Ordinary and durable-batch native subagents submit coroutines directly to one persistent isolated event loop. Gateway/embedded startup installs one process-wide async FIFO admission controller (default 3 running, bounded queue). Direct `create_deerflow_agent` callers can instead pass a caller-owned `SubagentRuntime`; reuse the same instance across graphs so its bound `task`, optional batch tools/service, middleware limits, and `SubagentExecutor` all share one controller without reading global YAML. An owned batch service must be started before graph construction and stopped at application shutdown. Waiters hold no scheduler thread, and cancellation/timeout release queue/slot ownership.
**Shared sandbox execution lifecycle** (#5128): every admitted subagent run carries a stable task-derived `sandbox_lease_owner_id` and matching `sandbox_command_scope_id` in its runtime context. Sandbox middleware retains that execution against the lead thread's active provider client, so one child finishing cannot close the sandbox while siblings still run; the final holder performs any pending provider release. A rollback/fork-restored child reusing the parent's live client binds a non-releasing holder: it fences parent cleanup and owns its command scope without requesting a park itself; a parent's earlier park request waits for the child, while a missing inherited client falls through to a normal fresh acquire. On AIO, the command scope selects one explicit persistent shell session per subagent, allowing independent scopes to run concurrently while preserving in-order shell state within one child. Sync sandbox tool bodies offloaded with `asyncio.to_thread` are shielded and drained across repeated cancellation before the outer execution can clean its holder; a cancelled worker can therefore neither re-admit an already-released owner nor run after subagent terminalization. Middleware performs the normal release, and `SubagentExecutor` repeats it idempotently in `finally` so exceptions, cooperative cancellation, and timeout unwind paths cannot leak a lease or scoped session.
**Shared sandbox execution lifecycle** (#5128): every admitted subagent run carries a stable task-derived `sandbox_lease_owner_id` and matching `sandbox_command_scope_id` in its runtime context. Sandbox middleware retains that execution against the lead thread's active provider client, so one child finishing cannot close the sandbox while siblings still run; the final holder performs any pending provider release. A rollback/fork-restored child reusing the parent's live client binds a non-releasing holder: it fences parent cleanup and owns its command scope without requesting a park itself; a parent's earlier park request waits for the child, while a missing inherited client falls through to a normal fresh acquire. On AIO, the command scope selects one explicit persistent shell session per subagent, allowing independent scopes to run concurrently while preserving in-order shell state within one child. Sync sandbox tool bodies offloaded with `asyncio.to_thread` are shielded and drained across repeated cancellation before the outer execution can clean its holder; a cancelled worker can therefore neither re-admit an already-released owner nor run after subagent terminalization. `SubagentExecutor` invokes active LangGraph stream cleanup before releasing its sandbox lease or notifying task stop. Slow cooperative cleanup keeps the result non-terminal and retains its sandbox lease and capacity slot; its warning is diagnostic, not a safe hard-timeout boundary. Middleware performs the normal release, and `SubagentExecutor` repeats it idempotently in `finally` so exceptions, cooperative cancellation, and timeout unwind paths cannot leak a lease or scoped session.
**Concurrency and total delegation cap**: Ordinary `task` concurrency is resolved once as the minimum of the per-run request, the startup-frozen `subagent_runtime.max_running`, and the schema safety ceiling (1-64), then shared by the lead prompt and `SubagentLimitMiddleware`. Hot reloads must not make either layer advertise more capacity than the already-created process controller; a changed startup-only value takes effect only after restart. The same middleware separately enforces `subagents.max_total_per_run` (default 6, config schema 1-50, runtime override `max_total_subagents` clamped to the same range) against current-run entries in the durable delegation ledger, so a long lead-agent run cannot bypass concurrency limits by launching repeated legal-sized batches at each planning checkpoint, but historical delegations from previous runs in the same thread do not consume the new run's budget. Explicit `batch_task` work does not consume or relax that ordinary-run ledger: its persisted total/live/running limits live under `subagent_batches`. Gateway `run_agent()` and embedded `DeerFlowClient.stream()` both provide a per-invocation `run_id` in runtime context; `DeerFlowClient.stream()` also tags its input `HumanMessage` with that same id so durable-context capture can identify the current request boundary. Gateway resume paths may not append a new `HumanMessage`, so the worker also exposes the pre-run checkpoint's message ids in runtime context; durable-context capture uses that as the current-run boundary and never re-tags older task calls as the resumed run. When no delegation slots remain, task calls are stripped, provider raw tool-call metadata is synced, `finish_reason` is forced to `stop`, and a visible "subagent delegation limit" note is appended so the agent can synthesize already-collected results. Default subagent timeout `subagents.timeout_seconds=1800` (30 min) and built-in `general-purpose` `max_turns=150`.
**Flow**: Ordinary `task()``SubagentExecutor` → shared process slot → result polling/SSE. Explicit `batch_task()` → durable batch/item rows → lease-based batch service (`subagents/batch_service.py`, started by Gateway or an explicit direct runtime) → the same `SubagentExecutor`/process slots → bounded stored result and owner-scoped API/JSONL export. Batch mode is selected only by the explicit tool, never inferred from prompt size. Executor queue rejection/timeout occurs before model execution and therefore releases the durable lease without consuming an item attempt; real execution failure and expired leases still consume the retry budget. User cancellation terminalizes every nonterminal item immediately and clears its lease, fencing any stale worker completion. Background cancellation resolves the result/future under `_background_tasks_lock` but calls `Future.cancel()` only after releasing it, because cancellation may synchronously invoke the completion callback that reacquires the registry lock. Direct runtimes provide the tools and worker but not Gateway's HTTP/UI surface. `task_started` carries the resolved effective model name. The per-subagent `SubagentTokenCollector` publishes a cumulative usage snapshot to the shared `SubagentResult` after every completed LLM response; the next `task_running` event carries that snapshot, so collapsed workspace cards can update without re-accounting parent-run totals. Terminal ToolMessage metadata (`subagent_model_name`, `subagent_token_usage`) and the persisted `subagent.end` event retain the model/usage after reload; absent provider usage stays absent rather than being estimated as zero. The executor caches one resolved `AppConfig` snapshot (explicit or `get_app_config()` fallback) for agent assembly, deferred setup, and receipt harvesting, so `verification.receipts_enabled=false` remains authoritative on both construction paths. Terminal tool receipts are harvested before `try_set_terminal` and committed with the other payload fields under the same state lock, so status polling cannot observe a terminal result before its receipt metadata is available. Each yielded values chunk becomes the latest terminal-harvest state and immediately publishes its harvested receipts to the shared result before cooperative cancellation is checked. Tool-ended cancellation/failure evidence uses the current ToolMessage scan, but a completed result always uses the bounded ledger snapshot attached to the assistant text being returned—even when a max-turn partial ends on a later tool chunk—so omitted receipts cannot validate its citations; a missing/malformed completed snapshot fails closed with no receipts. Therefore direct task cancellation and both execution/polling timeouts retain the latest execution evidence even when cancellation interrupts before another stream boundary.
**Report contract (RFC #4651 PR3)**: `report_contract.py` owns the prompt-layer text that makes Layer 1 receipt verification non-inert. `SubagentExecutor._build_initial_state` appends `build_report_contract_section(receipts_enabled=...)` to every subagent's consolidated `SystemMessage` — built-in and custom alike — requiring `[rN tool_name]` citations (from the Tool receipts ledger) for action claims, verifiable handles (absolute path, URL, ID, HTTP status) for deliverables, and explicit reporting of failures; the citation clause follows `verification.receipts_enabled`, and the citation example derives from the single-owner `format_citation`/`receipt_id` so prompt text cannot drift from the verifier. The `task` tool hands lead-supplied `acceptance_criteria` to the `SubagentExecutor` constructor, which appends them via `render_acceptance_criteria_block(...)` to the task `HumanMessage` (stripped, capped at 20 items × 500 chars, each entry neutralized) — the untrusted channel `InputSanitizationMiddleware` escapes and boundary-frames, matching their model-supplied provenance. The subagent's `SystemMessage` never carries criterion text; it gets only the framework-owned `build_acceptance_criteria_system_note(...)` pointer naming the list's location and authority, so natural-language injection inside a criterion cannot gain system-channel priority over framework instructions. Deterministic leaf checking is a separate layer.

View File

@ -6,6 +6,7 @@ import json
import logging
import os
import re
import sys
import threading
import uuid
from collections.abc import Callable, Coroutine, Mapping
@ -36,6 +37,7 @@ from deerflow.authz.principal import normalize_authz_attributes
from deerflow.config import get_app_config
from deerflow.config.app_config import AppConfig
from deerflow.models import create_chat_model
from deerflow.runtime.runs.stream_cleanup import close_agent_stream
from deerflow.runtime.user_context import DEFAULT_USER_ID
from deerflow.skills.types import Skill
from deerflow.subagents.capacity import (
@ -67,6 +69,7 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS = 3.0
_STREAM_CLOSE_SLOW_WARNING_SECONDS = 10.0
# Kept as wire keys here instead of importing ``deerflow.sandbox`` at module
# load: executor tests and extension embedders replace that package while
# breaking agent/tool import cycles.
@ -1587,41 +1590,77 @@ class SubagentExecutor:
)
return result
async for chunk in agent.astream(state, config=run_config, context=context, stream_mode="values"): # type: ignore[arg-type]
# A yielded values chunk is already executed state. Retain it
# before observing cooperative cancellation so terminal receipt
# harvesting includes a tool result that completed while the
# cancellation request was in flight.
final_state = chunk
result.update_tool_receipts(terminal_receipts())
result.update_bash_executions(current_bash_executions())
cancelled_during_stream = False
stream = agent.astream(state, config=run_config, context=context, stream_mode="values") # type: ignore[arg-type]
try:
async for chunk in stream:
# A yielded values chunk is already executed state. Retain it
# before observing cooperative cancellation so terminal receipt
# harvesting includes a tool result that completed while the
# cancellation request was in flight.
final_state = chunk
result.update_tool_receipts(terminal_receipts())
result.update_bash_executions(current_bash_executions())
# Cooperative cancellation: check if parent requested stop.
# Note: cancellation is only detected at astream iteration boundaries,
# so long-running tool calls within a single iteration will not be
# interrupted until the next chunk is yielded.
if result.cancel_event.is_set():
logger.info(f"[trace={self.trace_id}] Subagent {self.config.name} cancelled by parent")
result.try_set_terminal(
SubagentStatus.CANCELLED,
error="Cancelled by user",
token_usage_records=collector.snapshot_records(),
tool_receipts=terminal_receipts(),
# Cooperative cancellation: check if parent requested stop.
# Note: cancellation is only detected at astream iteration boundaries,
# so long-running tool calls within a single iteration will not be
# interrupted until the next chunk is yielded.
if result.cancel_event.is_set():
logger.info(f"[trace={self.trace_id}] Subagent {self.config.name} cancelled by parent")
cancelled_during_stream = True
break
result.update_token_usage_records(collector.snapshot_records())
# Capture every step message (assistant turns AND tool outputs)
# appended since the last chunk. A single super-step can append
# several ToolMessages when the model emits multiple tool calls in
# one turn, so capturing only messages[-1] would drop all but the
# last output (#3779). Dedup/serialization live in capture_step_message.
messages = chunk.get("messages", [])
previous_count = len(ai_messages)
processed_message_count = capture_new_step_messages(messages, ai_messages, seen_message_ids, processed_message_count)
if len(ai_messages) > previous_count:
logger.info(f"[trace={self.trace_id}] Subagent {self.config.name} captured {len(ai_messages) - previous_count} step message(s); total #{len(ai_messages)}")
finally:
active_error = sys.exception()
cancel_requested = cancelled_during_stream or result.cancel_event.is_set()
slow_close_warning = asyncio.get_running_loop().call_later(
_STREAM_CLOSE_SLOW_WARNING_SECONDS,
logger.warning,
"[trace=%s] Subagent %s stream cleanup for execution %s is still running after %.1fs; cooperative terminalization and sandbox resource release have not occurred",
self.trace_id,
self.config.name,
result.task_id,
_STREAM_CLOSE_SLOW_WARNING_SECONDS,
context=Context(),
)
try:
await close_agent_stream(stream)
except Exception:
cancel_requested = cancel_requested or result.cancel_event.is_set()
if active_error is None and not cancel_requested:
raise
logger.warning(
"[trace=%s] Could not close interrupted subagent stream %s",
self.trace_id,
self.config.name,
exc_info=True,
)
return result
else:
cancel_requested = cancel_requested or result.cancel_event.is_set()
finally:
slow_close_warning.cancel()
result.update_token_usage_records(collector.snapshot_records())
# Capture every step message (assistant turns AND tool outputs)
# appended since the last chunk. A single super-step can append
# several ToolMessages when the model emits multiple tool calls in
# one turn, so capturing only messages[-1] would drop all but the
# last output (#3779). Dedup/serialization live in capture_step_message.
messages = chunk.get("messages", [])
previous_count = len(ai_messages)
processed_message_count = capture_new_step_messages(messages, ai_messages, seen_message_ids, processed_message_count)
if len(ai_messages) > previous_count:
logger.info(f"[trace={self.trace_id}] Subagent {self.config.name} captured {len(ai_messages) - previous_count} step message(s); total #{len(ai_messages)}")
if cancel_requested:
result.try_set_terminal(
SubagentStatus.CANCELLED,
error="Cancelled by user",
token_usage_records=collector.snapshot_records(),
tool_receipts=terminal_receipts(),
)
return result
logger.info(f"[trace={self.trace_id}] Subagent {self.config.name} completed async execution")
token_usage_records = collector.snapshot_records()

View File

@ -20,6 +20,7 @@ import inspect
import sys
import threading
import time
from contextvars import Context
from datetime import datetime
from importlib.metadata import version as package_version
from pathlib import Path
@ -1352,7 +1353,26 @@ class TestAsyncExecutionPath:
final_message,
]
}
mock_agent.astream = lambda *args, **kwargs: async_iterator([final_state])
class Stream:
def __init__(self):
self.closed = False
self.yielded = False
def __aiter__(self):
return self
async def __anext__(self):
if self.yielded:
raise StopAsyncIteration
self.yielded = True
return final_state
async def aclose(self):
self.closed = True
stream = Stream()
mock_agent.astream.return_value = stream
executor = SubagentExecutor(
config=base_config,
@ -1369,6 +1389,7 @@ class TestAsyncExecutionPath:
assert result.error is None
assert result.started_at is not None
assert result.completed_at is not None
assert stream.closed is True
@pytest.mark.anyio
async def test_aexecute_marks_capacity_rejection_as_admission_failure(self, classes, base_config):
@ -1778,6 +1799,184 @@ class TestAsyncExecutionPath:
assert "Agent error" in result.error
assert result.completed_at is not None
@pytest.mark.anyio
async def test_aexecute_preserves_stream_error_when_close_also_fails(
self,
classes,
base_config,
mock_agent,
):
SubagentExecutor = classes["SubagentExecutor"]
SubagentStatus = classes["SubagentStatus"]
close_attempted = False
class Stream:
def __aiter__(self):
return self
async def __anext__(self):
raise ValueError("stream failed")
async def aclose(self):
nonlocal close_attempted
close_attempted = True
raise RuntimeError("close failed")
mock_agent.astream.return_value = Stream()
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")
assert result.status == SubagentStatus.FAILED
assert result.error == "stream failed"
assert close_attempted is True
@pytest.mark.anyio
async def test_aexecute_stream_error_wins_over_cooperative_cancellation(
self,
classes,
base_config,
mock_agent,
caplog,
):
SubagentExecutor = classes["SubagentExecutor"]
SubagentResult = classes["SubagentResult"]
SubagentStatus = classes["SubagentStatus"]
cancel_event = threading.Event()
class Stream:
def __aiter__(self):
return self
async def __anext__(self):
cancel_event.set()
raise ValueError("stream failed")
async def aclose(self):
return None
mock_agent.astream.return_value = Stream()
result_holder = SubagentResult(
task_id="cancel-with-stream-error",
trace_id="test-trace",
status=SubagentStatus.RUNNING,
started_at=datetime.now(),
)
result_holder.cancel_event = cancel_event
executor = SubagentExecutor(
config=base_config,
tools=[],
thread_id="test-thread",
)
with (
patch.object(executor, "_create_agent", return_value=mock_agent),
caplog.at_level("ERROR", logger="deerflow.subagents.executor"),
):
result = await executor._aexecute(
"Task",
result_holder=result_holder,
)
assert result.status is SubagentStatus.FAILED
assert result.error == "stream failed"
assert any(record.exc_info is not None and "async execution failed" in record.getMessage() for record in caplog.records)
@pytest.mark.anyio
async def test_aexecute_close_only_failure_marks_execution_failed(
self,
classes,
base_config,
mock_agent,
monkeypatch,
):
executor_module = importlib.import_module("deerflow.subagents.executor")
SubagentExecutor = classes["SubagentExecutor"]
SubagentStatus = classes["SubagentStatus"]
warning_handle = MagicMock()
loop = asyncio.get_running_loop()
real_call_later = loop.call_later
def record_warning_handle(delay, callback, *args, **kwargs):
if callback == executor_module.logger.warning:
assert delay == executor_module._STREAM_CLOSE_SLOW_WARNING_SECONDS
warning_context = kwargs.get("context")
assert isinstance(warning_context, Context)
assert list(warning_context.items()) == []
return warning_handle
return real_call_later(delay, callback, *args, **kwargs)
class Stream:
def __aiter__(self):
return self
async def __anext__(self):
raise StopAsyncIteration
async def aclose(self):
raise RuntimeError("close failed")
mock_agent.astream.return_value = Stream()
executor = SubagentExecutor(
config=base_config,
tools=[],
thread_id="test-thread",
)
monkeypatch.setattr(loop, "call_later", record_warning_handle)
with patch.object(executor, "_create_agent", return_value=mock_agent):
result = await executor._aexecute("Task")
assert result.status == SubagentStatus.FAILED
assert result.error == "close failed"
warning_handle.cancel.assert_called_once_with()
@pytest.mark.anyio
async def test_aexecute_close_error_does_not_mask_inflight_cancelled_error(
self,
classes,
base_config,
mock_agent,
):
SubagentExecutor = classes["SubagentExecutor"]
started = asyncio.Event()
close_attempted = False
class Stream:
def __aiter__(self):
return self
async def __anext__(self):
started.set()
await asyncio.Event().wait()
async def aclose(self):
nonlocal close_attempted
close_attempted = True
raise RuntimeError("close failed")
mock_agent.astream.return_value = Stream()
executor = SubagentExecutor(
config=base_config,
tools=[],
thread_id="test-thread",
)
with patch.object(executor, "_create_agent", return_value=mock_agent):
execution = asyncio.create_task(executor._aexecute("Task"))
await started.wait()
execution.cancel("host cancellation")
with pytest.raises(asyncio.CancelledError) as raised:
await execution
assert raised.value.args == ("host cancellation",)
assert close_attempted is True
@pytest.mark.anyio
async def test_aexecute_finally_releases_only_the_failing_subagent_lease(
self,
@ -3104,23 +3303,81 @@ class TestCooperativeCancellation:
assert call_count == 0 # astream was never entered
@pytest.mark.anyio
async def test_aexecute_cancelled_mid_stream(self, classes, base_config, msg):
@pytest.mark.parametrize("close_fails", [False, True])
async def test_aexecute_cancelled_mid_stream(
self,
classes,
base_config,
monkeypatch,
msg,
close_fails,
):
"""Test that _aexecute returns CANCELLED when cancel_event is set during streaming."""
SubagentExecutor = classes["SubagentExecutor"]
SubagentResult = classes["SubagentResult"]
SubagentStatus = classes["SubagentStatus"]
cancel_event = threading.Event()
events: list[str] = []
async def mock_astream(*args, **kwargs):
yield {"messages": [msg.human("Task"), msg.ai("Partial", "msg-1")]}
# Simulate cancellation during streaming
cancel_event.set()
yield {"messages": [msg.human("Task"), msg.ai("Should not appear", "msg-2")]}
class Stream:
def __init__(self):
self.yielded = 0
def __aiter__(self):
return self
async def __anext__(self):
self.yielded += 1
if self.yielded == 1:
return {
"messages": [
msg.human("Task"),
msg.ai("Partial", "msg-1"),
],
}
if self.yielded == 2:
cancel_event.set()
return {
"messages": [
msg.human("Task"),
msg.ai("Should not appear", "msg-2"),
],
}
raise StopAsyncIteration
async def aclose(self):
events.append("aclose")
if close_fails:
raise RuntimeError("close failed")
mock_agent = MagicMock()
stream = Stream()
def mock_astream(*_args, **kwargs):
kwargs["context"]["sandbox_id"] = "sandbox-1"
return stream
mock_agent.astream = mock_astream
class LeaseManager:
async def release_async(self, _owner_id):
events.append("lease_release")
sandbox_module = sys.modules["deerflow.sandbox"]
monkeypatch.setattr(
sandbox_module,
"get_sandbox_provider",
lambda: object(),
raising=False,
)
lease_module = importlib.import_module("deerflow.sandbox.lease")
monkeypatch.setattr(
lease_module,
"get_sandbox_lease_manager",
lambda _provider: LeaseManager(),
)
result_holder = SubagentResult(
task_id="cancel-mid",
trace_id="test-trace",
@ -3134,13 +3391,246 @@ class TestCooperativeCancellation:
tools=[],
thread_id="test-thread",
)
original_try_set_terminal = result_holder.try_set_terminal
with patch.object(executor, "_create_agent", return_value=mock_agent):
def record_terminal(*args, **kwargs):
if args[0] is SubagentStatus.CANCELLED:
events.append("terminal_cancelled")
return original_try_set_terminal(*args, **kwargs)
with (
patch.object(executor, "_create_agent", return_value=mock_agent),
patch.object(
result_holder,
"try_set_terminal",
side_effect=record_terminal,
),
):
result = await executor._aexecute("Task", result_holder=result_holder)
assert result.status == SubagentStatus.CANCELLED
assert result.error == "Cancelled by user"
assert result.completed_at is not None
assert events == [
"aclose",
"terminal_cancelled",
"lease_release",
]
@pytest.mark.anyio
@pytest.mark.parametrize("close_fails", [False, True])
async def test_aexecute_honors_cancellation_at_stream_exhaustion(
self,
classes,
base_config,
msg,
close_fails,
):
SubagentExecutor = classes["SubagentExecutor"]
SubagentResult = classes["SubagentResult"]
SubagentStatus = classes["SubagentStatus"]
cancel_event = threading.Event()
class Stream:
def __init__(self):
self.yielded = False
def __aiter__(self):
return self
async def __anext__(self):
if self.yielded:
cancel_event.set()
raise StopAsyncIteration
self.yielded = True
return {
"messages": [
msg.human("Task"),
msg.ai("Done", "msg-1"),
],
}
async def aclose(self):
if close_fails:
raise RuntimeError("close failed")
mock_agent = MagicMock()
mock_agent.astream.return_value = Stream()
result_holder = SubagentResult(
task_id="cancel-at-exhaustion",
trace_id="test-trace",
status=SubagentStatus.RUNNING,
started_at=datetime.now(),
)
result_holder.cancel_event = cancel_event
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",
result_holder=result_holder,
)
assert result.status == SubagentStatus.CANCELLED
assert result.error == "Cancelled by user"
@pytest.mark.anyio
@pytest.mark.parametrize("close_fails", [False, True])
async def test_aexecute_honors_cancellation_while_closing_stream(
self,
classes,
base_config,
msg,
close_fails,
):
SubagentExecutor = classes["SubagentExecutor"]
SubagentResult = classes["SubagentResult"]
SubagentStatus = classes["SubagentStatus"]
cancel_event = threading.Event()
close_started = asyncio.Event()
allow_close = asyncio.Event()
class Stream:
def __init__(self):
self.yielded = False
def __aiter__(self):
return self
async def __anext__(self):
if self.yielded:
raise StopAsyncIteration
self.yielded = True
return {
"messages": [
msg.human("Task"),
msg.ai("Done", "msg-1"),
],
}
async def aclose(self):
close_started.set()
await allow_close.wait()
if close_fails:
raise RuntimeError("close failed")
mock_agent = MagicMock()
mock_agent.astream.return_value = Stream()
result_holder = SubagentResult(
task_id="cancel-while-closing",
trace_id="test-trace",
status=SubagentStatus.RUNNING,
started_at=datetime.now(),
)
result_holder.cancel_event = cancel_event
executor = SubagentExecutor(
config=base_config,
tools=[],
thread_id="test-thread",
)
with patch.object(executor, "_create_agent", return_value=mock_agent):
execution = asyncio.create_task(
executor._aexecute(
"Task",
result_holder=result_holder,
)
)
await close_started.wait()
cancel_event.set()
allow_close.set()
result = await execution
assert result.status == SubagentStatus.CANCELLED
assert result.error == "Cancelled by user"
@pytest.mark.anyio
async def test_aexecute_warns_while_stream_cleanup_blocks_terminalization(
self,
classes,
base_config,
monkeypatch,
msg,
caplog,
):
executor_module = importlib.import_module("deerflow.subagents.executor")
SubagentExecutor = classes["SubagentExecutor"]
SubagentResult = classes["SubagentResult"]
SubagentStatus = classes["SubagentStatus"]
cancel_event = threading.Event()
close_started = asyncio.Event()
allow_close = asyncio.Event()
class Stream:
def __init__(self):
self.yielded = False
def __aiter__(self):
return self
async def __anext__(self):
if self.yielded:
raise StopAsyncIteration
self.yielded = True
cancel_event.set()
return {
"messages": [
msg.human("Task"),
msg.ai("Partial", "msg-1"),
],
}
async def aclose(self):
close_started.set()
await allow_close.wait()
mock_agent = MagicMock()
mock_agent.astream.return_value = Stream()
result_holder = SubagentResult(
task_id="slow-close",
trace_id="test-trace",
status=SubagentStatus.RUNNING,
started_at=datetime.now(),
)
result_holder.cancel_event = cancel_event
executor = SubagentExecutor(
config=base_config,
tools=[],
thread_id="test-thread",
)
monkeypatch.setattr(
executor_module,
"_STREAM_CLOSE_SLOW_WARNING_SECONDS",
0.0,
)
with (
patch.object(executor, "_create_agent", return_value=mock_agent),
caplog.at_level("WARNING", logger="deerflow.subagents.executor"),
):
execution = asyncio.create_task(
executor._aexecute(
"Task",
result_holder=result_holder,
)
)
await close_started.wait()
await asyncio.sleep(0)
assert not execution.done()
assert result_holder.status is SubagentStatus.RUNNING
assert [record for record in caplog.records if "stream cleanup for execution slow-close is still running" in record.getMessage()]
allow_close.set()
result = await execution
await asyncio.sleep(0)
assert result.status is SubagentStatus.CANCELLED
assert len([record for record in caplog.records if "stream cleanup for execution slow-close is still running" in record.getMessage()]) == 1
def test_request_cancel_sets_event(self, executor_module, classes):
"""Test that request_cancel_background_task sets the cancel_event."""