fix(subagents): drain final cleanup across caller cancellation (#7546)

* fix(subagents): drain final cleanup across caller cancellation

* fix(subagents): preserve stop outcomes during cancellation

Follow up on PR #7546 discussion_r4374235193 and Issue #7535.

Capture the unwinding cancellation before the drained cleanup tail and use it
only while the result remains non-terminal. Preserve completed, failed and
timed-out decisions. Keep cancellation propagation and cleanup ownership.

Share the synchronous result holder and publish its existing FAILED outcome
before requesting cancellation. This necessary timeout boundary overlaps #7483;
no active-cancel terminal write is added to the cleanup tail.

Validation: 12 targeted regressions pass on Windows and Linux; make test
37876 passed / 808 skipped, make test-blocking-io 292 passed / 1 skipped.
Format, lint, diff checks and four-file incremental OCR delegated review pass.

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
chaodidiao 2026-10-26 23:04:45 +08:00 committed by GitHub
parent 7a2dda1326
commit 6e4bcc16aa
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 242 additions and 35 deletions

View File

@ -2538,6 +2538,12 @@ Web chats accept direct model choices such as “ask MODEL to handle this”,
When a sub-agent ends with `return_direct=True` tools, including tools contributed by extension middleware, their outputs are returned in tool-call order. A failed tool marks the task as failed while preserving the batch outputs.
After its stream closes, a cancelled sub-agent keeps its execution slot until
sandbox cleanup and extension task-stop notifications finish under their existing
error and timeout policies. Repeated cancellation does not abandon this cleanup.
Task-stop reports an interrupted execution as aborted, preserving any completed,
failed or timed-out outcome already published before cancellation.
Ordinary `task` calls accept `context_mode="isolated"` (default) or
`context_mode="snapshot"`. Isolated tasks receive their delegated prompt as
before. Snapshot tasks also receive the parent's retained conversation and

View File

@ -1,5 +1,7 @@
### Subagent System (`packages/harness/deerflow/subagents/`)
Drain cleanup; preserve terminal outcomes on cancel.
Accumulate validated ToolMessage knowledge across values frames in
`SubagentResult.knowledge_sources`: at most 100 sources, deduplicated by evidence ID.
Publish before terminal state; retain through compaction. Ordinary `task` forwards

View File

@ -62,6 +62,7 @@ from deerflow.tracing import (
inject_langfuse_metadata,
)
from deerflow.utils.bash_exit import BASH_TOOL_NAMES, bash_exit_status
from deerflow.utils.file_io import await_drained
from deerflow.utils.messages import message_content_to_text
if TYPE_CHECKING:
@ -2095,37 +2096,42 @@ class SubagentExecutor:
# stream teardown, including cancellation and hard-timeout cleanup.
if collector is not None:
result.update_token_usage_records(collector.snapshot_records())
if execution_context is not None:
try:
from deerflow.sandbox.lease import release_sandbox_execution_lease_async
await release_sandbox_execution_lease_async(execution_context)
except Exception:
logger.warning(
"[trace=%s] Failed to release sandbox execution lease for subagent %s",
self.trace_id,
self.config.name,
exc_info=True,
)
if task_info is not None and task_store is not None:
try:
await notify_task_stop(
loaded_extensions,
task_store,
task_info,
subagent_task_outcome(
cancelled=result.status is SubagentStatus.CANCELLED,
succeeded=result.status is SubagentStatus.COMPLETED,
),
timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS,
)
except Exception:
logger.warning(
"[trace=%s] Extension task-stop notification failed for subagent %s (non-fatal)",
self.trace_id,
self.config.name,
exc_info=True,
)
async def cleanup_execution() -> None:
if execution_context is not None:
try:
from deerflow.sandbox.lease import release_sandbox_execution_lease_async
await release_sandbox_execution_lease_async(execution_context)
except Exception:
logger.warning(
"[trace=%s] Failed to release sandbox execution lease for subagent %s",
self.trace_id,
self.config.name,
exc_info=True,
)
if task_info is not None and task_store is not None:
try:
await notify_task_stop(
loaded_extensions,
task_store,
task_info,
subagent_task_outcome(
cancelled=result.status is SubagentStatus.CANCELLED,
succeeded=result.status is SubagentStatus.COMPLETED,
),
timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS,
)
except Exception:
logger.warning(
"[trace=%s] Extension task-stop notification failed for subagent %s (non-fatal)",
self.trace_id,
self.config.name,
exc_info=True,
)
# Keep capacity and execution-done ownership until the whole tail settles.
await await_drained(cleanup_execution())
return result
@ -2138,6 +2144,8 @@ class SubagentExecutor:
long-lived isolated loop. Reusing that loop keeps shared async clients
from being tied to a short-lived loop that gets closed per execution.
"""
if result_holder is None:
result_holder = SubagentResult(task_id=str(uuid.uuid4())[:8], trace_id=self.trace_id, status=SubagentStatus.PENDING)
future: Future[SubagentResult] | None = None
parent_context = _copy_isolated_subagent_context()
try:
@ -2147,11 +2155,10 @@ class SubagentExecutor:
)
return future.result(timeout=self.config.timeout_seconds)
except FuturesTimeoutError as exc:
if result_holder is not None:
# Publish the synchronous API's failure before cancellation can
# reach task-stop observers or win the terminal-state race.
result_holder.try_set_terminal(SubagentStatus.FAILED, error=str(exc))
result_holder.cancel_event.set()
# Publish the synchronous API's failure before cancellation can
# reach task-stop observers or win the terminal-state race.
result_holder.try_set_terminal(SubagentStatus.FAILED, error=str(exc))
result_holder.cancel_event.set()
if future is not None:
future.cancel()
raise

View File

@ -7016,6 +7016,198 @@ async def test_deferred_cleanup_expiry_registers_removal_after_fence_opens(class
executor_module.force_cleanup_background_task(execution_id)
class TestSubagentFinalizationCancellation:
@pytest.mark.anyio
@pytest.mark.parametrize("cleanup_stage", ["lease", "notify"])
@pytest.mark.parametrize(
"interrupted,terminal_status,expected_outcome",
[(False, None, "completed"), (True, None, "aborted"), (True, "TIMED_OUT", "failed"), (True, "FAILED", "failed")],
)
async def test_caller_cancellation_drains_execution_tail(self, classes, base_config, monkeypatch, cleanup_stage, interrupted, terminal_status, expected_outcome):
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
from deerflow.extensions import notify
from deerflow.extensions.registry import ExtensionRegistry
from deerflow.subagents.capacity import SubagentExecutionCapacity
loop = asyncio.get_running_loop()
monkeypatch.setattr(notify, "_notify_loop", loop)
stream_started, stream_closed, cleanup_started = asyncio.Event(), asyncio.Event(), asyncio.Event()
notification_release = asyncio.Event()
lease_release = threading.Event()
finished = []
outcomes = []
sandbox = MagicMock()
provider = MagicMock()
provider.get.return_value = sandbox
manager = SandboxLeaseManager(provider)
manager.retain("lead", "shared", thread_id="cleanup-thread", user_id="default")
def close_scope(owner_id):
if owner_id != "lead":
if cleanup_stage == "lease":
loop.call_soon_threadsafe(cleanup_started.set)
assert lease_release.wait(5), "test did not release sandbox cleanup"
finished.append("lease")
sandbox.release_command_scope.side_effect = close_scope
monkeypatch.setattr(importlib.import_module("deerflow.sandbox.sandbox_provider"), "get_sandbox_provider", lambda: provider)
monkeypatch.setattr(importlib.import_module("deerflow.sandbox.lease"), "get_sandbox_lease_manager", lambda _provider: manager)
class Observer:
async def on_task_start(self, *args):
finished.append("start")
async def on_task_stop(self, app_store, task_store, info, outcome):
if cleanup_stage == "notify":
cleanup_started.set()
await notification_release.wait()
outcomes.append(outcome.value)
finished.append("stop")
registry = ExtensionRegistry()
with registry.attributed_to("cleanup:test"):
registry.task_lifecycle(Observer())
class Agent:
async def astream(self, *args, context, **kwargs):
owner_id = context["sandbox_lease_owner_id"]
context["sandbox_id"] = "shared"
manager.retain(owner_id, "shared", thread_id="cleanup-thread", user_id="default")
try:
stream_started.set()
if interrupted:
await asyncio.Event().wait()
yield {"messages": [classes["AIMessage"](content="done")]}
finally:
stream_closed.set()
capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1))
executor = classes["SubagentExecutor"](config=base_config, tools=[], thread_id="cleanup-thread", run_id="cleanup-run", extensions=registry.build(), execution_capacity=capacity)
monkeypatch.setattr(executor, "_build_initial_state", AsyncMock(return_value=({}, [], None)))
monkeypatch.setattr(executor, "_create_agent", AsyncMock(return_value=Agent()))
holder = classes["SubagentResult"](task_id="tail-cleanup", trace_id="trace", status=classes["SubagentStatus"].PENDING)
execution = asyncio.create_task(executor._aexecute("Task", holder))
try:
if interrupted:
await asyncio.wait_for(stream_started.wait(), 5)
if terminal_status is not None:
holder.try_set_terminal(classes["SubagentStatus"][terminal_status], error="published before cancellation")
holder.cancel_event.set()
execution.cancel()
await asyncio.wait_for(cleanup_started.wait(), 5)
assert stream_closed.is_set()
for _ in range(3):
execution.cancel()
await asyncio.sleep(0)
assert not execution.done()
assert capacity.snapshot().running == 1
assert not holder.execution_done_event.is_set()
finally:
lease_release.set()
notification_release.set()
if not execution.done():
execution.cancel()
try:
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(execution, 5)
finally:
manager.release("subagent:tail-cleanup")
manager.release("lead")
assert finished == ["start", "lease", "stop"]
assert manager.binding_for("subagent:tail-cleanup") is None
assert capacity.snapshot().running == 0
assert holder.execution_done_event.is_set()
assert outcomes == [expected_outcome]
if terminal_status is not None:
assert holder.status is classes["SubagentStatus"][terminal_status]
if not interrupted:
assert holder.status is classes["SubagentStatus"].COMPLETED
@pytest.mark.parametrize("with_holder", [False, True])
def test_sync_timeout_keeps_failed_notification(self, classes, base_config, monkeypatch, with_holder):
from deerflow.extensions.registry import ExtensionRegistry
stopped, drained = threading.Event(), threading.Event()
outcomes = []
class Observer:
async def on_task_start(self, *args):
pass
async def on_task_stop(self, app_store, task_store, info, outcome):
outcomes.append(outcome.value)
stopped.set()
registry = ExtensionRegistry()
with registry.attributed_to("cleanup:test"):
registry.task_lifecycle(Observer())
class Agent:
async def astream(self, *args, **kwargs):
await asyncio.Event().wait()
yield {}
base_config.timeout_seconds = 0.2
executor = classes["SubagentExecutor"](config=base_config, tools=[], run_id="cleanup-run", extensions=registry.build())
monkeypatch.setattr(executor, "_build_initial_state", AsyncMock(return_value=({}, [], None)))
monkeypatch.setattr(executor, "_create_agent", AsyncMock(return_value=Agent()))
original_execute = executor._aexecute
async def record_drain(task, result_holder=None):
try:
return await original_execute(task, result_holder)
finally:
drained.set()
monkeypatch.setattr(executor, "_aexecute", record_drain)
holder = classes["SubagentResult"](task_id="sync-timeout", trace_id="trace", status=classes["SubagentStatus"].PENDING) if with_holder else None
result = executor.execute("Task", holder)
assert result.status is classes["SubagentStatus"].FAILED
assert stopped.wait(5), "timed-out execution did not notify observers"
assert drained.wait(5), "timed-out execution did not release its resources"
assert outcomes == ["failed"]
@pytest.mark.anyio
@pytest.mark.parametrize("failure", ["error", "timeout"])
async def test_task_stop_failure_keeps_existing_budget_and_result(self, classes, base_config, monkeypatch, failure):
from deerflow.extensions import notify
from deerflow.extensions.registry import ExtensionRegistry
monkeypatch.setattr(notify, "_notify_loop", asyncio.get_running_loop())
monkeypatch.setattr(importlib.import_module("deerflow.subagents.executor"), "_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS", 0.01)
entered, exited = asyncio.Event(), asyncio.Event()
class Observer:
async def on_task_start(self, *args):
pass
async def on_task_stop(self, *args):
entered.set()
try:
if failure == "error":
raise RuntimeError("observer failed")
await asyncio.Event().wait()
finally:
exited.set()
registry = ExtensionRegistry()
with registry.attributed_to("cleanup:test"):
registry.task_lifecycle(Observer())
class Agent:
async def astream(self, *args, **kwargs):
yield {"messages": [classes["AIMessage"](content="done")]}
executor = classes["SubagentExecutor"](config=base_config, tools=[], run_id="cleanup-run", extensions=registry.build())
monkeypatch.setattr(executor, "_build_initial_state", AsyncMock(return_value=({}, [], None)))
monkeypatch.setattr(executor, "_create_agent", AsyncMock(return_value=Agent()))
result = await executor._aexecute("Task")
assert entered.is_set() and exited.is_set()
assert result.status is classes["SubagentStatus"].COMPLETED
assert result.execution_done_event.is_set()
class TestSubagentCancellationLifecycle:
@pytest.fixture
def lifecycle(self):