mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
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:
parent
7a2dda1326
commit
6e4bcc16aa
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user