diff --git a/README.md b/README.md index 085887512e..422c561aeb 100644 --- a/README.md +++ b/README.md @@ -2212,6 +2212,9 @@ not wrapped by middleware model-call hooks (goal, memory, title, and summarizati Gateway-lifetime services, and eager FastAPI HTTP routers. The contract package has no framework dependencies; extensions must declare FastAPI, LangChain, LangGraph, or other libraries they import. +Cancelled subagents report `aborted` to task-lifecycle hooks after stream cleanup; +synchronous timeouts remain failures, and background hard timeouts retain their +failure outcome. The [fetched-content screening example](examples/deerflow-extension-jev-screening/README.md) contributes one such middleware at the visible tool position. After operator opt-in it classifies the sanitized text a remote tool result shows the model, PII-redacted when diff --git a/backend/packages/harness/deerflow/subagents/AGENTS.md b/backend/packages/harness/deerflow/subagents/AGENTS.md index 6ba9bd3945..edc4f14cc5 100644 --- a/backend/packages/harness/deerflow/subagents/AGENTS.md +++ b/backend/packages/harness/deerflow/subagents/AGENTS.md @@ -48,7 +48,10 @@ citation references, and opaque base64/URL/image sources. **Interrupted usage**: After stream cleanup, publish the collector's final snapshot on cancellation/timeout; callbacks can precede values chunks. -`execute_async()` owns terminalization; `SubagentResult` guards first write. +Cancellation terminalizes before task-stop notification and still propagates; +Sync `execute()` shares the worker's result and publishes FAILED before timeout +cancellation; `execute_async()` publishes TIMED_OUT. First terminal write wins. +Stream cleanup precedes notification; late cancellation cannot replace a timeout. Parent journals deduplicate by `source_run_id`; batch writes need a live lease. **Durable batch acceptance**: Normalize optional `batch_task` criteria before diff --git a/backend/packages/harness/deerflow/subagents/executor.py b/backend/packages/harness/deerflow/subagents/executor.py index 79c1973553..6f7414437b 100644 --- a/backend/packages/harness/deerflow/subagents/executor.py +++ b/backend/packages/harness/deerflow/subagents/executor.py @@ -2068,6 +2068,18 @@ class SubagentExecutor: tool_receipts=terminal_receipts(), ) + except asyncio.CancelledError: + # The outer background wrapper observes cancellation only after this + # function's task-stop hook. Publish it here so observers see ABORTED; + # a previously published hard timeout still wins the terminal race. + result.try_set_terminal( + SubagentStatus.CANCELLED, + error="Cancelled by user", + token_usage_records=collector.snapshot_records() if collector is not None else None, + tool_receipts=terminal_receipts(), + ) + raise + except Exception as e: logger.exception(f"[trace={self.trace_id}] Subagent {self.config.name} async execution failed") result.try_set_terminal( @@ -2078,11 +2090,9 @@ class SubagentExecutor: ) finally: - # Cancellation (including wait_for's execution timeout) bypasses - # the terminal branches above. A completed model response can be - # in the collector before its graph node publishes a values chunk. - # Stream teardown has drained at this point; publish its final - # usage before the outer wrapper chooses CANCELLED or TIMED_OUT. + # A completed model response can be in the collector before its + # graph node publishes a values chunk. Publish final usage after + # 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: @@ -2136,8 +2146,11 @@ class SubagentExecutor: lambda: self._aexecute(task, result_holder), ) return future.result(timeout=self.config.timeout_seconds) - except FuturesTimeoutError: + 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() if future is not None: future.cancel() @@ -2169,19 +2182,19 @@ class SubagentExecutor: Returns: SubagentResult with the execution result. """ + # Share one holder with the worker even when the caller did not supply + # one, so a timeout and its cleanup report the same terminal outcome. + result = result_holder + if result is None: + result = SubagentResult( + task_id=str(uuid.uuid4())[:8], + trace_id=self.trace_id, + status=SubagentStatus.PENDING, + ) try: - return self._execute_in_isolated_loop(task, result_holder) + return self._execute_in_isolated_loop(task, result) except Exception as e: logger.exception(f"[trace={self.trace_id}] Subagent {self.config.name} execution failed") - # Create a result with error if we don't have one - if result_holder is not None: - result = result_holder - else: - result = SubagentResult( - task_id=str(uuid.uuid4())[:8], - trace_id=self.trace_id, - status=SubagentStatus.RUNNING, - ) result.try_set_terminal(SubagentStatus.FAILED, error=str(e)) return result diff --git a/backend/tests/test_subagent_executor.py b/backend/tests/test_subagent_executor.py index c444ba7448..9e3ace08c8 100644 --- a/backend/tests/test_subagent_executor.py +++ b/backend/tests/test_subagent_executor.py @@ -7014,3 +7014,259 @@ async def test_deferred_cleanup_expiry_registers_removal_after_fence_opens(class assert executor_module.get_background_task_result(execution_id) is None finally: executor_module.force_cleanup_background_task(execution_id) + + +class TestSubagentCancellationLifecycle: + @pytest.fixture + def lifecycle(self): + from deerflow.extensions.registry import ExtensionRegistry + + starts, stops = [], [] + stopped = threading.Event() + + async def on_start(app_store, task_store, info): + starts.append(info) + + async def on_stop(app_store, task_store, info, outcome): + stops.append((info, outcome)) + stopped.set() + + registry = ExtensionRegistry() + with registry.attributed_to("test:lifecycle"): + registry.task_lifecycle(SimpleNamespace(on_task_start=on_start, on_task_stop=on_stop)) + return SimpleNamespace(extensions=registry.build(), starts=starts, stops=stops, stopped=stopped) + + @pytest.mark.parametrize("shared_holder", [False, True]) + @pytest.mark.parametrize("stage", ["setup", "stream"]) + def test_sync_timeout_preserves_failure_before_cancellation(self, classes, base_config, lifecycle, monkeypatch, shared_holder, stage): + from dataclasses import replace + + from deerflow_extension_api import TaskOutcome + + module = importlib.import_module("deerflow.subagents.executor") + entered, closed = threading.Event(), threading.Event() + worker_holders = [] + executor = classes["SubagentExecutor"](config=replace(base_config, timeout_seconds=0.2), tools=[], run_id="run", extensions=lifecycle.extensions) + holder = classes["SubagentResult"](task_id="sync-timeout", trace_id="trace", status=classes["SubagentStatus"].PENDING) if shared_holder else None + original_execute = executor._aexecute + original_submit = module._submit_to_isolated_loop_in_context + + async def initial_state(task): + if stage == "setup": + entered.set() + await asyncio.Event().wait() + return {}, [], None + + async def stream(*args, **kwargs): + try: + entered.set() + await asyncio.Event().wait() + yield {} + finally: + closed.set() + + async def track_worker(task, result_holder=None): + worker_holders.append(result_holder) + return await original_execute(task, result_holder) + + def submit(context, factory): + future = original_submit(context, factory) + original_cancel = future.cancel + + def cancel_and_drain(): + cancelled = original_cancel() + # Force cancellation handling to finish before execute() catches + # the timeout, so a shared-holder race cannot pass by scheduling luck. + assert lifecycle.stopped.wait(timeout=5) + return cancelled + + monkeypatch.setattr(future, "cancel", cancel_and_drain) + return future + + monkeypatch.setattr(executor, "_build_initial_state", initial_state) + monkeypatch.setattr(executor, "_create_agent", AsyncMock(return_value=SimpleNamespace(astream=stream))) + monkeypatch.setattr(executor, "_aexecute", track_worker) + monkeypatch.setattr(module, "_submit_to_isolated_loop_in_context", submit) + monkeypatch.setattr(module, "build_tracing_callbacks", lambda: []) + monkeypatch.setattr(module, "inject_langfuse_metadata", lambda *args, **kwargs: None) + try: + result = executor.execute("synthetic sync timeout", holder) + assert entered.is_set() + assert lifecycle.stopped.wait(timeout=5) + assert result.status is classes["SubagentStatus"].FAILED + assert result.error == "" # Preserve the synchronous TimeoutError result. + assert lifecycle.stops == [(lifecycle.starts[0], TaskOutcome.FAILED)] + assert worker_holders == [result] + assert result.task_id == lifecycle.starts[0].task_id + assert not shared_holder or result is holder + assert stage == "setup" or closed.is_set() + assert result.execution_done_event.wait(timeout=5) + finally: + module._shutdown_isolated_subagent_loop() + + @pytest.mark.parametrize("shared_holder", [False, True]) + def test_sync_timeout_returns_failure_before_stream_cleanup(self, classes, base_config, lifecycle, monkeypatch, shared_holder): + from dataclasses import replace + + from deerflow_extension_api import TaskOutcome + + from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig + from deerflow.subagents.capacity import SubagentExecutionCapacity + + module = importlib.import_module("deerflow.subagents.executor") + closing, release_close = threading.Event(), threading.Event() + capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1)) + + class Stream: + def __aiter__(self): + return self + + async def __anext__(self): + await asyncio.Event().wait() + + async def aclose(self): + closing.set() + while not release_close.is_set(): + await asyncio.sleep(0.001) + + executor = classes["SubagentExecutor"](config=replace(base_config, timeout_seconds=0.2), tools=[], run_id="run", extensions=lifecycle.extensions, execution_capacity=capacity) + holder = classes["SubagentResult"](task_id="slow-close", trace_id="trace", status=classes["SubagentStatus"].PENDING) if shared_holder else None + monkeypatch.setattr(executor, "_build_initial_state", AsyncMock(return_value=({}, [], None))) + monkeypatch.setattr(executor, "_create_agent", AsyncMock(return_value=SimpleNamespace(astream=lambda *args, **kwargs: Stream()))) + monkeypatch.setattr(module, "build_tracing_callbacks", lambda: []) + monkeypatch.setattr(module, "inject_langfuse_metadata", lambda *args, **kwargs: None) + try: + result = executor.execute("synthetic slow cleanup", holder) + assert closing.wait(timeout=5) + assert result.status is classes["SubagentStatus"].FAILED + assert result.error == "" + assert not shared_holder or result is holder + assert not result.execution_done_event.is_set() + assert capacity.snapshot().running == 1 + assert lifecycle.stops == [] + release_close.set() + assert result.execution_done_event.wait(timeout=5) + assert result.status is classes["SubagentStatus"].FAILED + assert capacity.snapshot().running == 0 + assert lifecycle.stops == [(lifecycle.starts[0], TaskOutcome.FAILED)] + finally: + release_close.set() + module._shutdown_isolated_subagent_loop() + + @pytest.mark.parametrize("mode", ["setup_cancel", "stream_cancel", "cooperative", "success", "failure", "timeout"]) + def test_background_outcome_matches_terminal_status(self, classes, base_config, lifecycle, monkeypatch, mode): + from dataclasses import replace + + from deerflow_extension_api import TaskOutcome + + from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig + from deerflow.subagents.capacity import SubagentExecutionCapacity + + module = importlib.import_module("deerflow.subagents.executor") + entered = threading.Event() + cooperative_cancel = threading.Event() + closed = threading.Event() + capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1)) + + async def initial_state(task): + if mode == "setup_cancel": + entered.set() + await asyncio.Event().wait() + return {}, [], None + + async def stream(*args, **kwargs): + try: + entered.set() + if mode in {"stream_cancel", "timeout"}: + await asyncio.Event().wait() + elif mode == "cooperative": + while not cooperative_cancel.is_set(): + await asyncio.sleep(0.001) + elif mode == "failure": + raise RuntimeError("synthetic execution failure") + yield {"messages": [classes["AIMessage"](content="done")]} + finally: + closed.set() + + config = replace(base_config, timeout_seconds=0.2 if mode == "timeout" else 10) + executor = classes["SubagentExecutor"](config=config, tools=[], thread_id="thread", run_id="run", extensions=lifecycle.extensions, execution_capacity=capacity) + monkeypatch.setattr(executor, "_build_initial_state", initial_state) + monkeypatch.setattr(executor, "_create_agent", AsyncMock(return_value=SimpleNamespace(astream=stream))) + monkeypatch.setattr(module, "build_tracing_callbacks", lambda: []) + monkeypatch.setattr(module, "inject_langfuse_metadata", lambda *args, **kwargs: None) + execution_id = executor.execute_async("synthetic lifecycle task") + result = module.get_background_task_result(execution_id) + try: + assert entered.wait(timeout=5) + if mode in {"setup_cancel", "stream_cancel"}: + module.request_cancel_background_task(execution_id) + elif mode == "cooperative": + result.cancel_event.set() + cooperative_cancel.set() + deadline = time.monotonic() + 5 + while not (result.status.is_terminal and result.execution_done_event.is_set()) and time.monotonic() < deadline: + time.sleep(0.001) + expected_status, expected_outcome = { + "success": (classes["SubagentStatus"].COMPLETED, TaskOutcome.COMPLETED), + "failure": (classes["SubagentStatus"].FAILED, TaskOutcome.FAILED), + "timeout": (classes["SubagentStatus"].TIMED_OUT, TaskOutcome.FAILED), + }.get(mode, (classes["SubagentStatus"].CANCELLED, TaskOutcome.ABORTED)) + assert result.status is expected_status + assert result.execution_done_event.is_set() + assert capacity.snapshot().running == 0 + assert mode == "setup_cancel" or closed.is_set() + assert len(lifecycle.starts) == 1 + assert lifecycle.stops == [(lifecycle.starts[0], expected_outcome)] + finally: + cooperative_cancel.set() + module.request_cancel_background_task(execution_id) + module._shutdown_isolated_subagent_loop() + module.cleanup_background_task(execution_id) + assert module.get_background_task_result(execution_id) is None + + @pytest.mark.anyio + async def test_direct_cancellation_drains_stream_before_notification(self, classes, base_config, lifecycle, monkeypatch): + from deerflow_extension_api import TaskOutcome + + from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig + from deerflow.subagents.capacity import SubagentExecutionCapacity + + entered, closing, release_close = asyncio.Event(), asyncio.Event(), asyncio.Event() + capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1)) + + class Stream: + def __aiter__(self): + return self + + async def __anext__(self): + entered.set() + await asyncio.Event().wait() + + async def aclose(self): + closing.set() + await release_close.wait() + + executor = classes["SubagentExecutor"](config=base_config, tools=[], run_id="run", extensions=lifecycle.extensions, execution_capacity=capacity) + monkeypatch.setattr(executor, "_build_initial_state", AsyncMock(return_value=({}, [], None))) + monkeypatch.setattr(executor, "_create_agent", AsyncMock(return_value=SimpleNamespace(astream=lambda *args, **kwargs: Stream()))) + holder = classes["SubagentResult"](task_id="cancel", trace_id="trace", status=classes["SubagentStatus"].PENDING) + task = asyncio.create_task(executor._aexecute("synthetic task", holder)) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + task.cancel() + await asyncio.wait_for(closing.wait(), timeout=5) + task.cancel() + await asyncio.sleep(0) + assert not task.done() + assert not holder.status.is_terminal + assert not holder.execution_done_event.is_set() + assert capacity.snapshot().running == 1 + assert lifecycle.stops == [] + finally: + release_close.set() + with pytest.raises(asyncio.CancelledError): + await task + assert holder.status is classes["SubagentStatus"].CANCELLED + assert holder.execution_done_event.is_set() + assert capacity.snapshot().running == 0 + assert lifecycle.stops == [(lifecycle.starts[0], TaskOutcome.ABORTED)]