diff --git a/backend/packages/harness/deerflow/subagents/AGENTS.md b/backend/packages/harness/deerflow/subagents/AGENTS.md index 65f3eacb0..84ecf8674 100644 --- a/backend/packages/harness/deerflow/subagents/AGENTS.md +++ b/backend/packages/harness/deerflow/subagents/AGENTS.md @@ -19,3 +19,5 @@ **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. **Isolated-loop callback boundary**: sync delegation from an active event loop and `execute_async()` copy the ambient ContextVars into the persistent subagent loop so checkpoint lineage, user identity, tracing context, tags, metadata, and LangGraph's namespaced message-stream handler survive. Before submission, `_copy_isolated_subagent_context()` copies the callback manager/list and removes only handlers marked `deerflow_loop_bound`; `RunJournal` carries that marker because it owns parent-loop tasks and a SQL store/pool. LangGraph merges inherited callbacks with the child run's explicit `SubagentTokenCollector`/tracing callbacks, so letting `RunJournal` cross loops causes duplicate accounting and `Future attached to a different loop` failures, while dropping the whole callback chain silently removes child token frames. Do not replace the boundary with a blank `Context`; the inherited checkpoint namespace and framework stream callback are required by the stream-isolation contract above. + +**Reverse direction of the loop boundary — deferred cleanup & final usage delivery (#5069)**: when a task-tool poller exits unexpectedly, the registry cleanup is pinned **to** the persistent subagent loop via the public `run_on_isolated_subagent_loop()` (executor) so it survives caller-loop teardown — `asyncio.run()` cancels caller-loop tasks on exit, so a caller-loop `asyncio.create_task` would be cancelled before running. The final usage report crosses the boundary the **other way**: `_schedule_deferred_subagent_cleanup` captures the parent run's loop at unwind time (alive in every path that continues the run — the polling-timeout branch returns normally, and a generic poller error becomes an error `ToolMessage`), and `_deliver_final_usage_report` hands the report back onto that loop with `call_soon_threadsafe`. `record_external_llm_usage_records` must never be invoked from the persistent loop or a worker thread (`to_thread`): the journal's accumulators are unlocked read-modify-write fields and `get_completion_data()` iterates `_tokens_by_model`, so a cross-thread write silently loses token updates or breaks iteration mid-run — calling `_report_subagent_usage` directly inside `_deferred_cleanup_subagent_task` would reintroduce exactly this race. The deferred cleaner captures only the resolved usage recorder (plus ids and the captured report loop) — never the whole `runtime`: the strongly-referenced cleanup task lives for up to the full poll budget, and through `runtime` it would pin the parent run's journal and event store for that entire window. When the captured parent loop is already closed (synchronous `asyncio.run` teardown), the report is dropped on purpose — the run has persisted its completion data and nothing reads the counters back — and logged at info with the execution id and unaccounted record count, because the registry entry is removed right after and those records exist nowhere else. diff --git a/backend/packages/harness/deerflow/subagents/executor.py b/backend/packages/harness/deerflow/subagents/executor.py index dfc822f3b..0aabda50f 100644 --- a/backend/packages/harness/deerflow/subagents/executor.py +++ b/backend/packages/harness/deerflow/subagents/executor.py @@ -470,6 +470,18 @@ def _submit_to_isolated_loop_in_context( return context.run(_submit) +def run_on_isolated_subagent_loop[T](coro: Coroutine[Any, Any, T]) -> Future[T]: + """Schedule a coroutine on the process-owned persistent subagent loop. + + Unlike ``asyncio.create_task`` on the caller's loop, work submitted here + survives teardown of a short-lived caller loop — e.g. the ``asyncio.run()`` + used by the synchronous tool wrapper cancels caller-loop tasks on exit — + so registry cleanup scheduled from a failing poller still runs after the + caller loop is gone. + """ + return asyncio.run_coroutine_threadsafe(coro, _get_isolated_subagent_loop()) + + def _copy_isolated_subagent_context() -> Context: """Copy ambient context without loop-bound parent graph callbacks. @@ -1639,3 +1651,24 @@ def cleanup_background_task(execution_id: str) -> None: execution_id, result.status.value if hasattr(result.status, "value") else result.status, ) + + +def force_cleanup_background_task(execution_id: str) -> None: + """Remove a background task entry unconditionally. + + Last resort for interrupted unwind paths where the registry entry exists + but its result object can no longer be read (persistent status-lookup / + status-object failure), so :func:`cleanup_background_task` — which reads + the entry to check terminality — cannot succeed. Cooperative cancellation + has already been requested by then; leaking the entry forever is worse + than dropping it. The subagent thread keeps its own reference to the + result object, so a later ``try_set_terminal`` on the removed object is + harmless. + + Args: + execution_id: The execution ID to remove. + """ + with _background_tasks_lock: + _background_tasks.pop(execution_id, None) + _background_futures.pop(execution_id, None) + logger.warning("Force-cleaned background execution %s after unreadable status", execution_id) diff --git a/backend/packages/harness/deerflow/tools/builtins/task_tool.py b/backend/packages/harness/deerflow/tools/builtins/task_tool.py index 15c770ab5..f8ad4d145 100644 --- a/backend/packages/harness/deerflow/tools/builtins/task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/task_tool.py @@ -1,7 +1,9 @@ """Task tool for delegating work to subagents.""" import asyncio +import concurrent.futures import logging +import time import uuid from contextvars import ContextVar from dataclasses import replace @@ -25,8 +27,10 @@ from deerflow.subagents.config import resolve_subagent_model_name from deerflow.subagents.executor import ( SubagentStatus, cleanup_background_task, + force_cleanup_background_task, get_background_task_result, request_cancel_background_task, + run_on_isolated_subagent_loop, ) from deerflow.subagents.status_contract import ( SubagentStatusValue, @@ -43,6 +47,23 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +# Poll cadence for terminal-state waits in both the interrupted unwind and the +# deferred registry cleaner. +_SUBAGENT_POLL_INTERVAL_SECONDS = 5.0 + +# How long the generic-error unwind waits for a terminal result before +# re-raising. This is deliberately a short grace period, not the full +# ``max_poll_count`` budget: a subagent blocked inside a long model/tool call +# may not observe cooperative cancellation for the whole execution timeout +# (~31 minutes by default), and an unrelated poller failure must not stall +# the parent run that long. The remaining lifecycle is handed to the deferred +# cleaner on the persistent subagent loop. +_UNEXPECTED_EXIT_GRACE_SECONDS = 5.0 + +# Sentinel returned by ``_peek_subagent_result`` when the registry entry exists +# but cannot be read (persistent status-lookup / status-object failure). +_STATUS_UNREADABLE = object() + _explicit_execution_capacity: ContextVar[SubagentExecutionCapacity | None] = ContextVar( "deerflow_explicit_subagent_execution_capacity", default=None, @@ -58,36 +79,257 @@ def _is_subagent_terminal(result: Any) -> bool: return result.status in {SubagentStatus.COMPLETED, SubagentStatus.FAILED, SubagentStatus.CANCELLED, SubagentStatus.TIMED_OUT} or getattr(result, "completed_at", None) is not None -async def _await_subagent_terminal(execution_id: str, max_polls: int) -> Any | None: - """Poll until the background subagent reaches a terminal status or we run out of polls.""" - for _ in range(max_polls): +def _peek_subagent_result(execution_id: str, *, trace_id: str) -> Any: + """Read a registry entry without letting a broken status object raise. + + The generic-error unwind exists to handle poller failures caused by + persistent status-lookup/status-object errors; finalization re-reading + through the same failing accessor must not abort the unwind. Returns the + entry when readable, ``None`` when it is gone (nothing left to clean), + and ``_STATUS_UNREADABLE`` when it exists but cannot be read. + """ + try: result = get_background_task_result(execution_id) - if result is None: - return None + except Exception: + logger.warning( + f"[trace={trace_id}] Background status lookup failed for execution {execution_id}", + exc_info=True, + ) + return _STATUS_UNREADABLE + if result is None: + return None + try: + _is_subagent_terminal(result) + except Exception: + logger.warning( + f"[trace={trace_id}] Background status object unreadable for execution {execution_id}", + exc_info=True, + ) + return _STATUS_UNREADABLE + return result + + +async def _await_subagent_terminal(execution_id: str, max_polls: int, *, trace_id: str = "", grace_seconds: float | None = None) -> Any: + """Poll until the background subagent reaches a terminal status. + + Without ``grace_seconds`` the wait is bounded by ``max_polls`` polls (the + cancellation unwind's contract). With it, the wait is additionally bounded + by wall-clock time — the generic-error unwind must re-raise promptly + instead of stalling the parent run for the full execution timeout. Never + raises through a broken status accessor; propagates ``_STATUS_UNREADABLE`` + to the caller instead. + """ + polls = 0 + deadline = None if grace_seconds is None else time.monotonic() + grace_seconds + while True: + result = _peek_subagent_result(execution_id, trace_id=trace_id) + if result is None or result is _STATUS_UNREADABLE: + return result if _is_subagent_terminal(result): return result - await asyncio.sleep(5) - return None + if polls >= max_polls - 1: + return None + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + return None + await asyncio.sleep(min(_SUBAGENT_POLL_INTERVAL_SECONDS, remaining)) + else: + await asyncio.sleep(_SUBAGENT_POLL_INTERVAL_SECONDS) + polls += 1 -async def _deferred_cleanup_subagent_task(execution_id: str, trace_id: str, max_polls: int) -> None: - """Keep polling a cancelled subagent until it can be safely removed.""" +async def _finalize_interrupted_subagent( + runtime: Runtime, + execution_id: str, + trace_id: str, + max_polls: int, + *, + grace_seconds: float | None = None, +) -> None: + """Shared unwind for interrupted polling (cancellation or unexpected error). + + Wait (shielded, bounded by ``max_polls`` or ``grace_seconds``) for the + subagent to reach a terminal state so the final token usage snapshot is + reported to the parent RunJournal, then remove the registry entry. + Terminal results are removed synchronously before re-raising; non-terminal + and unreadable ones defer removal to the process-owned persistent subagent + loop, which survives teardown of a short-lived caller loop (``asyncio.run`` + cancels caller-loop tasks — including any detached cleanup task — on exit). + + This function must never raise: it runs while an exception (the original + poller failure or cancellation) is already in flight, and any error it + raised would replace that exception and skip the cleanup attachment. A + persistently unreadable status object therefore falls through to the + deferred cleaner rather than propagating. + """ + try: + unreadable = False + terminal_result = None + try: + waited = await asyncio.shield(_await_subagent_terminal(execution_id, max_polls, trace_id=trace_id, grace_seconds=grace_seconds)) + if waited is _STATUS_UNREADABLE: + unreadable = True + else: + terminal_result = waited + except asyncio.CancelledError: + # The shielded wait surfaces an outer cancellation here. The + # cancellation branch REQUIRES this absorb — re-raising would + # abort the unwind before the deferred-cleanup attachment. The + # generic-error branch shares this helper, so a cancellation + # landing inside its grace wait is absorbed too; that branch + # re-checks task.cancelling() after the unwind and re-raises + # CancelledError so the node still ends as an interrupted run + # (see the unwind call site in task_tool). + pass + + # Report whatever the subagent collected (even if we timed out). + final_result = terminal_result + final_terminal = False + if final_result is not None: + final_terminal = _is_subagent_terminal(final_result) + else: + peek = _peek_subagent_result(execution_id, trace_id=trace_id) + if peek is _STATUS_UNREADABLE: + unreadable = True + elif peek is not None: + final_result = peek + final_terminal = _is_subagent_terminal(peek) + + if unreadable: + # The entry exists but cannot be read; the terminal-gated sync + # cleanup cannot be trusted here. Attach the deferred cleaner, + # whose last resort force-removes unreadable entries. + _schedule_deferred_subagent_cleanup(runtime, execution_id, trace_id, max_polls) + return + + if final_result is not None: + _report_subagent_usage(runtime, final_result) + if final_terminal: + cleanup_background_task(execution_id) + else: + _schedule_deferred_subagent_cleanup(runtime, execution_id, trace_id, max_polls) + except Exception: + logger.error( + f"[trace={trace_id}] Interrupted-subagent finalization failed for execution {execution_id}", + exc_info=True, + ) + + +def _deliver_final_usage_report( + usage_recorder: Any, + result: Any, + report_loop: asyncio.AbstractEventLoop | None, + *, + execution_id: str, +) -> None: + """Schedule the FINAL usage report onto the loop that owns the RunJournal. + + ``RunJournal`` is deliberately ``deerflow_loop_bound``: its accumulators + are unlocked read-modify-write fields and ``_tokens_by_model`` is iterated + by ``get_completion_data()``, so reporting from any other thread races the + parent run's own journal writes (lost token updates, ``dictionary changed + size during iteration``). ``report_loop`` is captured at unwind time, when + the unwind paths still run on the parent run's loop. That loop is alive in + every path that continues the run — the polling-timeout branch returns + normally and a generic poller error becomes an error ``ToolMessage``, both + handing control back to the lead agent — so ``call_soon_threadsafe`` + delivers the report on the journal's own loop, serialized with every other + journal access. ``usage_recorder`` is likewise resolved at unwind time: + the deferred cleaner must retain only the handler, not the whole + ``runtime`` (whose journal and event store belong to the parent run and + would be pinned for the cleaner's whole poll budget otherwise). + + A ``None`` recorder means this run has no journal at all — skip without + touching any loop. + + On the synchronous ``asyncio.run`` path the loop may already be closed by + the time the deferred cleaner reaches a terminal result. The report is + then dropped on purpose: the run has finished and persisted its completion + data, so nothing reads those counters back — recording into a dead run's + journal would account nothing. This is the one path where a subagent's + tail usage goes permanently unaccounted (the registry entry is removed + right after, so the records exist nowhere else), so the drop is logged at + info with the execution id and the record count. The report bypasses the + snapshot's ``usage_reported`` flag so records accumulated after the + snapshot are counted; the journal itself dedupes by ``source_run_id``. + """ + if usage_recorder is None: + logger.debug("Deferred final usage report for execution %s skipped: no usage recorder on this run", execution_id) + return + if report_loop is None: + logger.info( + "Dropping deferred final usage report for execution %s: no parent loop captured (%d usage records unaccounted)", + execution_id, + len(getattr(result, "token_usage_records", None) or []), + ) + return + try: + # A lambda, not plain arguments: call_soon_threadsafe forwards keyword + # arguments to the loop machinery (only ``context`` is its own), so + # passing ``final=True`` through it raises TypeError. + report_loop.call_soon_threadsafe(lambda: _report_usage_records(usage_recorder, result, final=True)) + except Exception: + # Loop closed between the capture and this call, or scheduling was + # rejected — same drop rationale and same info-level visibility. + # Never-raise by contract: delivery problems must not block the + # caller's registry removal. + logger.info( + "Dropping deferred final usage report for execution %s: parent loop closed before delivery (%d usage records unaccounted)", + execution_id, + len(getattr(result, "token_usage_records", None) or []), + ) + + +async def _deferred_cleanup_subagent_task( + usage_recorder: Any, + execution_id: str, + trace_id: str, + max_polls: int, + *, + report_loop: asyncio.AbstractEventLoop | None = None, +) -> None: + """Keep polling an interrupted subagent until it can be safely removed. + + Only the resolved usage recorder is retained (plus ids and the captured + report loop) — never the whole ``runtime``: the strongly-referenced + cleanup task lives for up to the full poll budget, and through + ``runtime`` it would pin the parent run's journal and event store for + that entire window. + + On a terminal result, schedule the subagent's FINAL usage report (deltas + since the unwind snapshot included) onto the parent run's loop BEFORE + removing the entry, so the parent RunJournal sees everything the subagent + collected — the scheduled callback holds its own reference to the result, + so removal does not invalidate the pending report. When the entry exists + but stays unreadable through the whole poll budget, force-remove it: + cooperative cancellation was already requested, and a broken status + object must not leak the entry forever. + """ cleanup_poll_count = 0 while True: - result = get_background_task_result(execution_id) + result = _peek_subagent_result(execution_id, trace_id=trace_id) if result is None: return - if _is_subagent_terminal(result): + if result is _STATUS_UNREADABLE: + if cleanup_poll_count >= max_polls: + logger.warning(f"[trace={trace_id}] Deferred cleanup for execution {execution_id}: status stayed unreadable after {cleanup_poll_count} polls, force-removing") + force_cleanup_background_task(execution_id) + return + elif _is_subagent_terminal(result): + # Never-raise by contract: delivery problems (closed parent loop) + # are logged inside and must not block the registry removal. + _deliver_final_usage_report(usage_recorder, result, report_loop, execution_id=execution_id) cleanup_background_task(execution_id) return if cleanup_poll_count >= max_polls: logger.warning(f"[trace={trace_id}] Deferred cleanup for execution {execution_id} timed out after {cleanup_poll_count} polls") return - await asyncio.sleep(5) + await asyncio.sleep(_SUBAGENT_POLL_INTERVAL_SECONDS) cleanup_poll_count += 1 -def _log_cleanup_failure(cleanup_task: asyncio.Task[None], *, trace_id: str, execution_id: str) -> None: +def _log_cleanup_failure(cleanup_task: asyncio.Task[None] | concurrent.futures.Future, *, trace_id: str, execution_id: str) -> None: if cleanup_task.cancelled(): return @@ -96,7 +338,11 @@ def _log_cleanup_failure(cleanup_task: asyncio.Task[None], *, trace_id: str, exe logger.error(f"[trace={trace_id}] Deferred cleanup failed for execution {execution_id}: {exc}") -_deferred_cleanup_tasks: set[asyncio.Task[None]] = set() +# Strong references to scheduled deferred cleanups. The event loop only keeps +# weak references to tasks, so an unreferenced cleanup could be garbage +# collected mid-poll; entries hold either an asyncio task (caller-loop +# fallback) or a concurrent future (persistent subagent loop). +_deferred_cleanup_tasks: set[asyncio.Task[None] | concurrent.futures.Future] = set() def bind_task_tool( @@ -128,13 +374,54 @@ def bind_task_tool( return task_tool.model_copy(update={"coroutine": bound_coroutine}) -def _schedule_deferred_subagent_cleanup(execution_id: str, trace_id: str, max_polls: int) -> asyncio.Task[None]: +def _schedule_deferred_subagent_cleanup( + runtime: Runtime, + execution_id: str, + trace_id: str, + max_polls: int, +) -> asyncio.Task[None] | concurrent.futures.Future: + """Schedule the deferred registry cleanup on the process-owned subagent loop. + + The persistent loop outlives the poller's own event loop, so the cleanup + still runs when the poller exits under synchronous tool invocation, where + ``asyncio.run()`` cancels caller-loop tasks at teardown before a detached + ``asyncio.create_task`` could execute. If the persistent loop cannot be + obtained, fall back to the caller loop rather than raising out of an + unwind path that is already handling an error. + """ logger.debug(f"[trace={trace_id}] Scheduling deferred cleanup for cancelled execution {execution_id}") - cleanup_task = asyncio.create_task(_deferred_cleanup_subagent_task(execution_id, trace_id, max_polls)) - _deferred_cleanup_tasks.add(cleanup_task) - cleanup_task.add_done_callback(_deferred_cleanup_tasks.discard) - cleanup_task.add_done_callback(lambda task: _log_cleanup_failure(task, trace_id=trace_id, execution_id=execution_id)) - return cleanup_task + # Resolve both cross-loop dependencies here, on the unwind path's loop: + # the parent run's loop, so the deferred final usage report can be + # delivered back onto the loop that owns the RunJournal (see + # ``_deliver_final_usage_report``), and the usage recorder itself, so the + # cleaner retains only the handler instead of pinning the whole + # ``runtime`` (journal + event store) for its whole poll budget. + try: + report_loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() + except RuntimeError: + report_loop = None + usage_recorder = _find_usage_recorder(runtime) + coro = _deferred_cleanup_subagent_task(usage_recorder, execution_id, trace_id, max_polls, report_loop=report_loop) + try: + cleanup_handle = run_on_isolated_subagent_loop(coro) + except Exception: + # Unreachable in practice — the persistent loop backs the subagent + # execution itself, so it exists by the time a poller needs cleanup. + logger.warning( + f"[trace={trace_id}] Persistent subagent loop unavailable for deferred cleanup of {execution_id}; falling back to the caller loop", + exc_info=True, + ) + try: + cleanup_handle = asyncio.create_task(coro) + except Exception: + # No caller loop either — close the coroutine so it is not left + # un-awaited, and let the unwind's error handling take over. + coro.close() + raise + _deferred_cleanup_tasks.add(cleanup_handle) + cleanup_handle.add_done_callback(_deferred_cleanup_tasks.discard) + cleanup_handle.add_done_callback(lambda task: _log_cleanup_failure(task, trace_id=trace_id, execution_id=execution_id)) + return cleanup_handle def _find_usage_recorder(runtime: Any) -> Any | None: @@ -180,27 +467,45 @@ def _summarize_usage(records: list[dict] | None) -> dict | None: } -def _report_subagent_usage(runtime: Any, result: Any) -> None: - """Report subagent token usage to the parent RunJournal, if available. +def _report_usage_records(recorder: Any, result: Any, *, final: bool = False) -> None: + """Deliver usage records to a resolved recorder (flag-gated, never raises). - Each subagent task must be reported only once (guarded by usage_reported). + Shared core of both report paths: the unwind reports directly with a + runtime (resolving the recorder on the parent loop), while the deferred + cleaner delivers onto the parent loop with the recorder resolved at + unwind time — retaining only the handler, never the whole ``runtime`` + (which pins the run's journal and event store for the cleaner's whole + poll budget otherwise). """ - if getattr(result, "usage_reported", True): + if not final and getattr(result, "usage_reported", True): return records = getattr(result, "token_usage_records", None) or [] if not records: return - journal = _find_usage_recorder(runtime) - if journal is None: + if recorder is None: logger.debug("No usage recorder found in runtime callbacks — subagent token usage not recorded") return try: - journal.record_external_llm_usage_records(records) + recorder.record_external_llm_usage_records(records) result.usage_reported = True except Exception: logger.warning("Failed to report subagent token usage", exc_info=True) +def _report_subagent_usage(runtime: Any, result: Any, *, final: bool = False) -> None: + """Report subagent token usage to the parent RunJournal, if available. + + Each subagent task's snapshot must be reported only once (guarded by + usage_reported). The deferred cleaner's final report bypasses that flag + via ``final=True``: records accumulated after the snapshot are still + delivered, and the journal dedupes per ``source_run_id`` so nothing is + double-counted. Both call sites run on the parent run's loop — directly + from the poller, or via ``call_soon_threadsafe`` from the deferred + cleaner — preserving the journal's ``deerflow_loop_bound`` contract. + """ + _report_usage_records(_find_usage_recorder(runtime), result, final=final) + + def _get_runtime_app_config(runtime: Any) -> "AppConfig | None": explicit = _explicit_app_config.get() if explicit is not None: @@ -513,18 +818,21 @@ async def task_tool( logger.info(f"[trace={trace_id}] Started background task {tool_call_id} (execution_id={execution_id}, subagent={subagent_type}, timeout={config.timeout_seconds}s, polling_limit={max_poll_count} polls)") writer = get_stream_writer() - # Send Task Started message' - await aemit_custom_event( - { - "type": "task_started", - "task_id": tool_call_id, - "description": description, - "model_name": effective_model, - }, - writer=writer, - ) - try: + # Send Task Started message. This is a real await point (registered + # handlers run here), so it belongs inside the guarded region: an emit + # failure must take the same cooperative-cancel + deferred-cleanup + # path as any other unexpected exit, not leak the background entry. + await aemit_custom_event( + { + "type": "task_started", + "task_id": tool_call_id, + "description": description, + "model_name": effective_model, + }, + writer=writer, + ) + while True: result = get_background_task_result(execution_id) @@ -707,7 +1015,7 @@ async def task_tool( # cancellation and schedule deferred cleanup to remove the entry from # _background_tasks once the background thread reaches a terminal state. request_cancel_background_task(execution_id) - _schedule_deferred_subagent_cleanup(execution_id, trace_id, max_poll_count) + _schedule_deferred_subagent_cleanup(runtime, execution_id, trace_id, max_poll_count) message = f"Task polling timed out after {timeout_minutes} minutes. This may indicate the background task is stuck. Status: {result.status.value}" return _task_result_command( tool_call_id=tool_call_id, @@ -718,24 +1026,44 @@ async def task_tool( tool_receipts=getattr(result, "tool_receipts", None), ) except asyncio.CancelledError: - # Signal the background subagent thread to stop cooperatively. - request_cancel_background_task(execution_id) - - # Wait (shielded) for the subagent to reach a terminal state so the - # final token usage snapshot is reported to the parent RunJournal - # before the parent worker persists get_completion_data(). - terminal_result = None + # Signal the background subagent thread to stop cooperatively, then + # wait for the terminal result so the final token usage snapshot is + # reported to the parent RunJournal before the parent worker persists + # get_completion_data(). A failure here must not replace the + # CancelledError that is already in flight. try: - terminal_result = await asyncio.shield(_await_subagent_terminal(execution_id, max_poll_count)) - except asyncio.CancelledError: - pass - - # Report whatever the subagent collected (even if we timed out). - final_result = terminal_result or get_background_task_result(execution_id) - if final_result is not None: - _report_subagent_usage(runtime, final_result) - if final_result is not None and _is_subagent_terminal(final_result): - cleanup_background_task(execution_id) - else: - _schedule_deferred_subagent_cleanup(execution_id, trace_id, max_poll_count) + request_cancel_background_task(execution_id) + except Exception: + logger.warning( + f"[trace={trace_id}] Failed to request cancellation for background task {execution_id} during unwind", + exc_info=True, + ) + await _finalize_interrupted_subagent(runtime, execution_id, trace_id, max_poll_count) + raise + except Exception: + # Unexpected poller failure (emit error, status-lookup bug, writer + # failure, ...). Mirror the cancellation unwind: stop the subagent + # cooperatively, report its final usage, and remove the registry entry — + # synchronously when it already reached terminal, otherwise via + # deferred cleanup pinned to the process-owned subagent loop so it + # survives asyncio.run() teardown on the synchronous tool path. The + # unwind is bounded by a short grace period (not the full execution + # timeout) and never lets a failing status accessor or cancellation + # request replace the original exception. + try: + request_cancel_background_task(execution_id) + except Exception: + logger.warning( + f"[trace={trace_id}] Failed to request cancellation for background task {execution_id} during unwind", + exc_info=True, + ) + await _finalize_interrupted_subagent(runtime, execution_id, trace_id, max_poll_count, grace_seconds=_UNEXPECTED_EXIT_GRACE_SECONDS) + current_task = asyncio.current_task() + if current_task is not None and current_task.cancelling(): + # A graph-node cancellation landed inside the grace wait and was + # absorbed by the shared unwind (its never-raise contract catches + # CancelledError so the deferred-cleanup attachment still runs). + # Honour it here rather than reporting a tool failure: the node + # must end as an interrupted run, not a failed tool call. + raise asyncio.CancelledError raise diff --git a/backend/tests/test_subagent_executor.py b/backend/tests/test_subagent_executor.py index 5fa968c51..478a44b13 100644 --- a/backend/tests/test_subagent_executor.py +++ b/backend/tests/test_subagent_executor.py @@ -2158,6 +2158,32 @@ class TestThreadSafety: return _patch_default_get_app_config(importlib.reload(executor)) + def test_run_on_isolated_subagent_loop_survives_caller_loop_teardown(self, executor_module): + """Pinning work to the process-owned persistent subagent loop must keep + it runnable after the short-lived caller loop is torn down. + + Deferred registry cleanup scheduled from a failing task-tool poller + relies on this: ``asyncio.run()`` cancels caller-loop tasks on exit, + so a cleanup submitted with ``asyncio.create_task`` on the caller's + loop dies at teardown, while one submitted through + ``run_on_isolated_subagent_loop`` still runs to completion.""" + completed = threading.Event() + handles = [] + + async def deferred_work() -> None: + completed.set() + + async def schedule_from_caller() -> None: + handles.append(executor_module.run_on_isolated_subagent_loop(deferred_work())) + + # asyncio.run() creates the caller loop, runs the scheduling, then + # closes the loop — cancelling anything still pending on it. + asyncio.run(schedule_from_caller()) + + assert completed.wait(timeout=10), "work pinned to the persistent subagent loop must run after caller-loop teardown" + assert handles[0].done() + assert handles[0].result(timeout=10) is None + def test_multiple_executors_in_parallel(self, classes, base_config, msg): """Test multiple executors running in parallel via thread pool.""" from concurrent.futures import ThreadPoolExecutor, as_completed @@ -3043,6 +3069,33 @@ class TestCooperativeCancellation: assert task_id not in executor_module._background_tasks + def test_force_cleanup_removes_unreadable_running_task(self, executor_module, classes): + """Force cleanup removes a RUNNING entry unconditionally. + + Last resort for interrupted unwinds where the status object can no + longer be read (persistent accessor failure), so the terminality + check inside cleanup_background_task cannot be trusted: cooperative + cancellation was already requested by the caller. + """ + SubagentResult = classes["SubagentResult"] + SubagentStatus = classes["SubagentStatus"] + + task_id = "test-force-cleanup-running" + result = SubagentResult( + task_id=task_id, + trace_id="test-trace", + status=SubagentStatus.RUNNING, + ) + executor_module._background_tasks[task_id] = result + + executor_module.force_cleanup_background_task(task_id) + + assert task_id not in executor_module._background_tasks + + def test_force_cleanup_handles_unknown_task_gracefully(self, executor_module): + """Force cleanup doesn't raise for unknown task IDs.""" + executor_module.force_cleanup_background_task("nonexistent-task") + # ----------------------------------------------------------------------------- # Subagent Tracing Wiring diff --git a/backend/tests/test_task_tool_core_logic.py b/backend/tests/test_task_tool_core_logic.py index 99ee94b6d..d67c4af63 100644 --- a/backend/tests/test_task_tool_core_logic.py +++ b/backend/tests/test_task_tool_core_logic.py @@ -4,6 +4,8 @@ import asyncio import gc import importlib import inspect +import threading +import time import weakref from enum import Enum from types import SimpleNamespace @@ -29,6 +31,8 @@ from deerflow.subagents.status_contract import ( ) # Use module import so tests can patch the exact symbols referenced inside task_tool(). +# NOTE: conftest.py replaces deerflow.subagents.executor with a MagicMock, so the +# executor-bound names inside task_tool are mocks; tests patch them explicitly. task_tool_module = importlib.import_module("deerflow.tools.builtins.task_tool") @@ -1369,7 +1373,7 @@ def test_cleanup_not_called_on_polling_safety_timeout(monkeypatch): def add_done_callback(self, _callback): return None - def fake_create_task(coro): + def fake_run_on_isolated_subagent_loop(coro): scheduled_cleanups.append(coro) coro.close() return DummyCleanupTask() @@ -1389,7 +1393,7 @@ def test_cleanup_not_called_on_polling_safety_timeout(monkeypatch): ) monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: events.append) monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep) - monkeypatch.setattr(task_tool_module.asyncio, "create_task", fake_create_task) + monkeypatch.setattr(task_tool_module, "run_on_isolated_subagent_loop", fake_run_on_isolated_subagent_loop) monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: []) monkeypatch.setattr( task_tool_module, @@ -1475,6 +1479,761 @@ def test_cleanup_scheduled_on_cancellation(monkeypatch): assert cleanup_calls == ["tc-cancelled-cleanup"] +def test_task_started_emit_failure_stops_subagent_reports_usage_and_cleans_up(monkeypatch): + """An exception from the task_started emit — a real await point before the + polling loop — must mirror the cancellation unwind: cooperative cancel, + final usage reported to the parent RunJournal, and synchronous registry + cleanup. A detached cleanup task would not survive asyncio.run() teardown + on the synchronous tool path, so the terminal case must clean up directly.""" + config = _make_subagent_config() + cancel_calls: list[str] = [] + cleanup_calls: list[str] = [] + reported: list = [] + terminal_result = _make_result(FakeSubagentStatus.COMPLETED, result="done") + + async def failing_emit(event, *, writer=None): + raise RuntimeError("emit boom") + + monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) + monkeypatch.setattr( + task_tool_module, + "SubagentExecutor", + type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: task_id}), + ) + monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) + monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _: terminal_result) + monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: lambda _event: None) + monkeypatch.setattr(task_tool_module, "aemit_custom_event", failing_emit) + monkeypatch.setattr(task_tool_module, "request_cancel_background_task", lambda execution_id: cancel_calls.append(execution_id)) + monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda execution_id: cleanup_calls.append(execution_id)) + monkeypatch.setattr(task_tool_module, "_report_subagent_usage", lambda runtime, result: reported.append(result)) + monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep) + monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: []) + + with pytest.raises(RuntimeError, match="emit boom"): + _run_task_tool( + runtime=_make_runtime(), + description="test", + prompt="p", + subagent_type="general-purpose", + tool_call_id="tc-emit-fail", + ) + + assert cancel_calls == ["tc-emit-fail"], "emit failure must cooperatively cancel the subagent" + assert reported == [terminal_result], "emit failure must still report the subagent's final usage" + assert cleanup_calls == ["tc-emit-fail"], "terminal subagent must be cleaned up synchronously" + + +def test_unexpected_poller_error_deferred_cleanup_survives_sync_invocation(monkeypatch): + """Non-terminal fallback under synchronous tool invocation: the poller + dies while the subagent never reaches terminal within the bounded wait, so + removal is deferred — and the deferred cleanup must actually run after + ``asyncio.run()`` tears down the caller loop. It runs on a loop that + outlives the poller (the process-owned persistent subagent loop in + production; an equivalent long-lived loop here), never on a caller-loop + task, which teardown cancels. The unwind, scheduling wrapper, and cleanup + coroutine are the real production code paths.""" + config = _make_subagent_config() + reported: list = [] + cancel_calls: list[str] = [] + cleanup_calls: list[str] = [] + emit_calls = 0 + scheduled_handles: list = [] + main_thread = threading.current_thread() + caller_loop_finished = threading.Event() + execution_id = "exec-sync-poller-death" + result = SimpleNamespace( + status=FakeSubagentStatus.RUNNING, + ai_messages=["partial"], + result=None, + ) + + persistent_loop = asyncio.new_event_loop() + loop_thread = threading.Thread(target=persistent_loop.run_forever, name="test-persistent-cleanup-loop", daemon=True) + loop_thread.start() + + async def fail_on_status_emit(event, *, writer=None): + nonlocal emit_calls + emit_calls += 1 + if emit_calls >= 2: + raise RuntimeError("status emit boom") + return None + + def flip_terminal_off_caller_thread(queried_id): + # The poller, its bounded unwind wait, and the final snapshot all run + # on the caller thread (asyncio.run) and must keep seeing RUNNING so + # the deferred path is taken; once the deferred cleaner polls from the + # long-lived loop thread, the subagent reaches terminal. Gated on the + # caller loop actually closing so the loop-pinned final report is + # deterministically dropped here (while that loop is alive, delivery + # is legitimate — that case is pinned by the live-loop test). + if threading.current_thread() is not main_thread: + caller_loop_finished.wait(timeout=10.0) + result.status = FakeSubagentStatus.COMPLETED + return result + + def transport_to_persistent_loop(coro): + # Same primitive production uses via run_on_isolated_subagent_loop: + # pin the coroutine to a loop that outlives the caller's asyncio.run() + # loop instead of a caller-loop asyncio.create_task. + handle = asyncio.run_coroutine_threadsafe(coro, persistent_loop) + scheduled_handles.append(handle) + return handle + + monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) + monkeypatch.setattr( + task_tool_module, + "SubagentExecutor", + type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: execution_id}), + ) + monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) + monkeypatch.setattr(task_tool_module, "get_background_task_result", flip_terminal_off_caller_thread) + monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: lambda _event: None) + monkeypatch.setattr(task_tool_module, "aemit_custom_event", fail_on_status_emit) + monkeypatch.setattr(task_tool_module, "request_cancel_background_task", lambda execution_id_arg: cancel_calls.append(execution_id_arg)) + monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda execution_id_arg: cleanup_calls.append(execution_id_arg)) + monkeypatch.setattr(task_tool_module, "run_on_isolated_subagent_loop", transport_to_persistent_loop) + monkeypatch.setattr(task_tool_module, "_report_subagent_usage", lambda runtime, r, **kwargs: reported.append(r)) + monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep) + monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: []) + + try: + with pytest.raises(RuntimeError, match="status emit boom"): + _run_task_tool( + runtime=_make_runtime(), + description="test", + prompt="p", + subagent_type="general-purpose", + tool_call_id="tc-poll-fail-sync", + ) + # asyncio.run() has returned: the caller loop is fully closed now, so + # the deferred cleaner may safely flip to terminal. + caller_loop_finished.set() + + # The caller loop is torn down by asyncio.run() at this point. The + # cooperative cancel already fired and the snapshot usage was + # reported on the caller loop itself... The stub accepts **kwargs so + # the deferred final report (final=True) would ALSO be recorded if it + # were ever delivered — the caller loop is closed by the time the + # deferred cleaner reaches terminal, so the loop-pinned delivery + # intentionally drops it and the count stays at exactly one. + assert cancel_calls == [execution_id], "unexpected poller exit must cooperatively cancel the subagent" + assert reported == [result], "snapshot usage reported on the caller loop; closed-loop final report must be dropped, not threaded in" + + # ...and the deferred cleanup, pinned to the long-lived loop, removes + # the registry entry even though the caller loop that scheduled it is + # gone. A caller-loop asyncio.create_task would have been cancelled at + # teardown and this would time out. + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline and not cleanup_calls: + time.sleep(0.05) + assert cleanup_calls == [execution_id], "deferred cleanup must run after caller-loop teardown" + assert result.status is FakeSubagentStatus.COMPLETED + finally: + for handle in scheduled_handles: + try: + handle.result(timeout=5) + except Exception: + pass + task_tool_module._deferred_cleanup_tasks.discard(handle) + persistent_loop.call_soon_threadsafe(persistent_loop.stop) + loop_thread.join(timeout=5) + persistent_loop.close() + + +def test_unexpected_error_with_failing_status_accessor_preserves_exception_and_attaches_cleanup(monkeypatch): + """Persistent status-lookup failure: the registry accessor itself raises + during polling and would raise again during finalization. The unwind must + still (1) cooperatively cancel, (2) attach cleanup through a mechanism + that does not depend on the failing accessor — the deferred cleaner, whose + last resort force-removes the unreadable entry — and (3) re-raise the + ORIGINAL poller exception, never one raised by finalization.""" + config = _make_subagent_config() + cancel_calls: list[str] = [] + cleanup_calls: list[str] = [] + force_cleanup_calls: list[str] = [] + reported: list = [] + emit_calls = 0 + scheduled_handles: list = [] + execution_id = "exec-status-accessor-broken" + + persistent_loop = asyncio.new_event_loop() + loop_thread = threading.Thread(target=persistent_loop.run_forever, name="test-broken-status-cleanup-loop", daemon=True) + loop_thread.start() + + async def ok_emit(event, *, writer=None): + nonlocal emit_calls + emit_calls += 1 + return None + + def broken_accessor(queried_id): + raise RuntimeError("registry lookup boom") + + def transport_to_persistent_loop(coro): + handle = asyncio.run_coroutine_threadsafe(coro, persistent_loop) + scheduled_handles.append(handle) + return handle + + monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) + monkeypatch.setattr( + task_tool_module, + "SubagentExecutor", + type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: execution_id}), + ) + monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) + monkeypatch.setattr(task_tool_module, "get_background_task_result", broken_accessor) + monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: lambda _event: None) + monkeypatch.setattr(task_tool_module, "aemit_custom_event", ok_emit) + monkeypatch.setattr(task_tool_module, "request_cancel_background_task", lambda execution_id_arg: cancel_calls.append(execution_id_arg)) + monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda execution_id_arg: cleanup_calls.append(execution_id_arg)) + monkeypatch.setattr(task_tool_module, "force_cleanup_background_task", lambda execution_id_arg: force_cleanup_calls.append(execution_id_arg)) + monkeypatch.setattr(task_tool_module, "run_on_isolated_subagent_loop", transport_to_persistent_loop) + monkeypatch.setattr(task_tool_module, "_report_subagent_usage", lambda runtime, r, **kwargs: reported.append(r)) + monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep) + monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: []) + + try: + # The ORIGINAL exception is the accessor failure from the polling + # loop; finalization re-reading through the same broken accessor must + # not replace it with anything else. + with pytest.raises(RuntimeError, match="registry lookup boom"): + _run_task_tool( + runtime=_make_runtime(), + description="test", + prompt="p", + subagent_type="general-purpose", + tool_call_id=execution_id, + ) + + assert cancel_calls == [execution_id], "broken status accessor must not prevent cooperative cancellation" + assert reported == [], "an unreadable result must not be force-reported" + + # The deferred cleaner keeps hitting the broken accessor and, once its + # poll budget is exhausted, force-removes the entry instead of leaking + # it forever. + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline and not force_cleanup_calls: + time.sleep(0.05) + assert force_cleanup_calls == [execution_id], "deferred cleaner must force-remove an unreadable registry entry" + assert cleanup_calls == [], "terminal-gated cleanup must not run for an unreadable entry" + finally: + for handle in scheduled_handles: + try: + handle.result(timeout=5) + except Exception: + pass + task_tool_module._deferred_cleanup_tasks.discard(handle) + persistent_loop.call_soon_threadsafe(persistent_loop.stop) + loop_thread.join(timeout=5) + persistent_loop.close() + + +def test_unexpected_error_re_raised_promptly_via_short_grace(monkeypatch): + """A generic poller failure must not stall the parent run for the full + execution timeout (~31 minutes by default) waiting on a subagent that may + never observe cooperative cancellation. The unwind waits only a short + grace period, then hands the remaining lifecycle to the deferred cleaner + and re-raises the original error promptly.""" + config = _make_subagent_config() + cancel_calls: list[str] = [] + scheduled_cleanups: list = [] + emit_calls = 0 + started = time.monotonic() + + class DummyCleanupTask: + def add_done_callback(self, _callback): + return None + + def fake_run_on_isolated_subagent_loop(coro): + scheduled_cleanups.append(coro) + coro.close() + return DummyCleanupTask() + + async def fail_on_second_emit(event, *, writer=None): + nonlocal emit_calls + emit_calls += 1 + if emit_calls >= 2: + raise RuntimeError("status emit boom") + return None + + monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) + monkeypatch.setattr( + task_tool_module, + "SubagentExecutor", + type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: "exec-grace"}), + ) + monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) + monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _: _make_result(FakeSubagentStatus.RUNNING, ai_messages=["partial"])) + monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: lambda _event: None) + monkeypatch.setattr(task_tool_module, "aemit_custom_event", fail_on_second_emit) + monkeypatch.setattr(task_tool_module, "request_cancel_background_task", lambda execution_id: cancel_calls.append(execution_id)) + monkeypatch.setattr(task_tool_module, "run_on_isolated_subagent_loop", fake_run_on_isolated_subagent_loop) + # Real asyncio.sleep + a tiny grace window: the unwind must return within + # the grace bound, not after max_poll_count * 5s. + monkeypatch.setattr(task_tool_module, "_UNEXPECTED_EXIT_GRACE_SECONDS", 0.2) + monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: []) + + with pytest.raises(RuntimeError, match="status emit boom"): + _run_task_tool( + runtime=_make_runtime(), + description="test", + prompt="p", + subagent_type="general-purpose", + tool_call_id="tc-grace", + ) + + elapsed = time.monotonic() - started + assert elapsed < 5.0, f"unexpected-error unwind must re-raise promptly, took {elapsed:.1f}s" + assert cancel_calls == ["exec-grace"], "grace-bounded unwind must still request cooperative cancellation" + assert len(scheduled_cleanups) == 1, "remaining lifecycle must be handed to the deferred cleaner" + + +def test_deferred_cleanup_drops_final_usage_when_parent_loop_closed(monkeypatch): + """The deferred cleaner removes the terminal entry, and the loop-pinned + final usage report is DROPPED — not threaded in — when the parent loop + that was captured at unwind time is already closed. + + On the synchronous ``asyncio.run`` path the run has finished and persisted + its completion data by the time the deferred cleaner reaches a terminal + result, so recording into the dead run's journal would account nothing; + a ``to_thread`` report would instead race the journal from a foreign + thread while some other run's loop may still touch it.""" + config = _make_subagent_config() + cleanup_calls: list[str] = [] + cancel_calls: list[str] = [] + emit_calls = 0 + scheduled_handles: list = [] + main_thread = threading.current_thread() + caller_loop_finished = threading.Event() + execution_id = "exec-deferred-final-usage" + result = SimpleNamespace( + status=FakeSubagentStatus.RUNNING, + ai_messages=["partial"], + result=None, + # The unwind's snapshot report already ran and set this flag; the + # deferred final report must bypass it rather than return early. + usage_reported=True, + token_usage_records=[{"source_run_id": "run-1", "total_tokens": 10}], + ) + + class LoopPinnedJournal: + """Real-recorder path: captures the running loop of every call, so a + cross-thread report surfaces as a wrong-loop (or no-loop) entry.""" + + def __init__(self) -> None: + self.calls: list[tuple[object, list]] = [] + + def record_external_llm_usage_records(self, records): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + self.calls.append((loop, list(records))) + + journal = LoopPinnedJournal() + + persistent_loop = asyncio.new_event_loop() + loop_thread = threading.Thread(target=persistent_loop.run_forever, name="test-final-usage-cleanup-loop", daemon=True) + loop_thread.start() + + async def fail_on_status_emit(event, *, writer=None): + nonlocal emit_calls + emit_calls += 1 + if emit_calls >= 2: + raise RuntimeError("status emit boom") + return None + + def flip_terminal_and_grow_usage_after_caller_teardown(queried_id): + if threading.current_thread() is not main_thread: + # Deterministic ordering: only reach terminal once the caller's + # asyncio.run() loop has fully returned and closed. Flipping any + # earlier would let the loop-pinned report legitimately deliver + # while that loop is still alive (correct, but not what this + # test pins down). + caller_loop_finished.wait(timeout=10.0) + result.status = FakeSubagentStatus.COMPLETED + # Usage accumulated after the snapshot the unwind reported. + result.token_usage_records = [ + {"source_run_id": "run-1", "total_tokens": 10}, + {"source_run_id": "run-2", "total_tokens": 25}, + ] + return result + + def transport_to_persistent_loop(coro): + handle = asyncio.run_coroutine_threadsafe(coro, persistent_loop) + scheduled_handles.append(handle) + return handle + + monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) + monkeypatch.setattr( + task_tool_module, + "SubagentExecutor", + type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: execution_id}), + ) + monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) + monkeypatch.setattr(task_tool_module, "get_background_task_result", flip_terminal_and_grow_usage_after_caller_teardown) + monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: lambda _event: None) + monkeypatch.setattr(task_tool_module, "aemit_custom_event", fail_on_status_emit) + monkeypatch.setattr(task_tool_module, "request_cancel_background_task", lambda execution_id_arg: cancel_calls.append(execution_id_arg)) + monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda execution_id_arg: cleanup_calls.append(execution_id_arg)) + monkeypatch.setattr(task_tool_module, "run_on_isolated_subagent_loop", transport_to_persistent_loop) + monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep) + monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: []) + + runtime = SimpleNamespace( + state={ + "sandbox": {"sandbox_id": "local"}, + "thread_data": { + "workspace_path": "/tmp/workspace", + "uploads_path": "/tmp/uploads", + "outputs_path": "/tmp/outputs", + }, + }, + context={"thread_id": "thread-1"}, + config={ + "metadata": {"model_name": "ark-model", "trace_id": "trace-1"}, + "callbacks": [journal], + }, + ) + + try: + with pytest.raises(RuntimeError, match="status emit boom"): + _run_task_tool( + runtime=runtime, + description="test", + prompt="p", + subagent_type="general-purpose", + tool_call_id="tc-final-usage", + ) + # asyncio.run() has returned: the caller loop is fully closed now, so + # the deferred cleaner may safely flip to terminal. + caller_loop_finished.set() + + assert cancel_calls == [execution_id] + + # The deferred cleaner reaches terminal off the caller thread and + # removes the entry. The caller loop is closed by then, so the + # loop-pinned final report is dropped: the recorder is real and + # present (resolved at unwind time), so an empty journal proves the + # drop is loop-based — a worker-thread report would have recorded + # with no running loop instead. + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline and not cleanup_calls: + time.sleep(0.05) + assert cleanup_calls == [execution_id], "deferred cleanup must remove the terminal entry" + assert journal.calls == [], "closed parent loop must drop the final report, not deliver it cross-thread" + finally: + for handle in scheduled_handles: + try: + handle.result(timeout=5) + except Exception: + pass + task_tool_module._deferred_cleanup_tasks.discard(handle) + persistent_loop.call_soon_threadsafe(persistent_loop.stop) + loop_thread.join(timeout=5) + persistent_loop.close() + + +def test_deferred_final_usage_reported_on_parent_loop_with_real_recorder(monkeypatch): + """Gateway shape: the parent run loop stays alive after the unexpected + poller exit, and the deferred final usage report must be DELIVERED onto + that loop — never from the deferred cleaner's thread. + + This exercises the real recorder path (``_find_usage_recorder`` → + ``_report_subagent_usage`` → ``journal.record_external_llm_usage_records``) + instead of stubbing ``_report_subagent_usage``: the journal captures the + running loop of every call, so a cross-thread report would surface here as + a wrong-loop (or no-loop) entry — the exact hazard ``deerflow_loop_bound`` + exists to prevent.""" + config = _make_subagent_config() + cleanup_calls: list[str] = [] + cancel_calls: list[str] = [] + emit_calls = 0 + scheduled_handles: list = [] + execution_id = "exec-final-usage-live-loop" + result = SimpleNamespace( + status=FakeSubagentStatus.RUNNING, + ai_messages=["partial"], + result=None, + # The snapshot report already ran during the run; the final report + # must bypass this flag and deliver the post-snapshot delta. + usage_reported=True, + token_usage_records=[{"source_run_id": "run-1", "total_tokens": 10}], + ) + + class LoopPinnedJournal: + def __init__(self) -> None: + self.calls: list[tuple[object, list]] = [] + + def record_external_llm_usage_records(self, records): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + self.calls.append((loop, list(records))) + + journal = LoopPinnedJournal() + + parent_loop = asyncio.new_event_loop() + parent_thread = threading.Thread(target=parent_loop.run_forever, name="test-live-parent-run-loop", daemon=True) + parent_thread.start() + persistent_loop = asyncio.new_event_loop() + persistent_thread = threading.Thread(target=persistent_loop.run_forever, name="test-final-usage-persistent-loop", daemon=True) + persistent_thread.start() + + async def fail_on_status_emit(event, *, writer=None): + nonlocal emit_calls + emit_calls += 1 + if emit_calls >= 2: + raise RuntimeError("status emit boom") + return None + + def flip_terminal_and_grow_usage_from_deferred_thread(queried_id): + if threading.current_thread() is persistent_thread: + result.status = FakeSubagentStatus.COMPLETED + # Usage accumulated after the snapshot the unwind reported. + result.token_usage_records = [ + {"source_run_id": "run-1", "total_tokens": 10}, + {"source_run_id": "run-2", "total_tokens": 25}, + ] + return result + + def transport_to_persistent_loop(coro): + handle = asyncio.run_coroutine_threadsafe(coro, persistent_loop) + scheduled_handles.append(handle) + return handle + + monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) + monkeypatch.setattr( + task_tool_module, + "SubagentExecutor", + type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: execution_id}), + ) + monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) + monkeypatch.setattr(task_tool_module, "get_background_task_result", flip_terminal_and_grow_usage_from_deferred_thread) + monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: lambda _event: None) + monkeypatch.setattr(task_tool_module, "aemit_custom_event", fail_on_status_emit) + monkeypatch.setattr(task_tool_module, "request_cancel_background_task", lambda execution_id_arg: cancel_calls.append(execution_id_arg)) + monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda execution_id_arg: cleanup_calls.append(execution_id_arg)) + monkeypatch.setattr(task_tool_module, "run_on_isolated_subagent_loop", transport_to_persistent_loop) + monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: []) + + runtime = SimpleNamespace( + state={ + "sandbox": {"sandbox_id": "local"}, + "thread_data": { + "workspace_path": "/tmp/workspace", + "uploads_path": "/tmp/uploads", + "outputs_path": "/tmp/outputs", + }, + }, + context={"thread_id": "thread-1"}, + config={ + "metadata": {"model_name": "ark-model", "trace_id": "trace-1"}, + "callbacks": [journal], + }, + ) + + try: + coroutine = getattr(task_tool_module.task_tool, "coroutine", None) + assert coroutine is not None + tool_future = asyncio.run_coroutine_threadsafe( + coroutine( + runtime=runtime, + description="test", + prompt="p", + subagent_type="general-purpose", + tool_call_id="tc-final-usage-live", + ), + parent_loop, + ) + with pytest.raises(RuntimeError, match="status emit boom"): + tool_future.result(timeout=10) + assert cancel_calls == [execution_id] + + # The deferred cleaner (persistent loop) reaches terminal, pins the + # final report onto the still-live parent loop, and removes the entry. + # Wait for the report to be OBSERVED on the parent loop, not merely + # scheduled — the journal only records once the callback actually ran. + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline and not journal.calls: + time.sleep(0.05) + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline and not cleanup_calls: + time.sleep(0.05) + assert cleanup_calls == [execution_id], "deferred cleanup must remove the terminal entry" + + assert len(journal.calls) == 1, "exactly one final usage report must be delivered" + observed_loop, observed_records = journal.calls[0] + assert observed_loop is parent_loop, f"record_external_llm_usage_records must run on the parent run loop that owns the RunJournal — got {observed_loop!r}" + assert [r["source_run_id"] for r in observed_records] == ["run-1", "run-2"], "the final report must include the post-snapshot delta records" + assert result.usage_reported is True + finally: + for handle in scheduled_handles: + try: + handle.result(timeout=5) + except Exception: + pass + task_tool_module._deferred_cleanup_tasks.discard(handle) + persistent_loop.call_soon_threadsafe(persistent_loop.stop) + persistent_thread.join(timeout=5) + persistent_loop.close() + parent_loop.call_soon_threadsafe(parent_loop.stop) + parent_thread.join(timeout=5) + parent_loop.close() + + +@pytest.mark.asyncio +async def test_unexpected_error_grace_wait_cancellation_is_honored(monkeypatch): + """A graph-node cancellation landing inside the generic-error grace wait + must surface as CancelledError, not the original poller error. + + The shared unwind absorbs CancelledError (its never-raise contract keeps + the deferred-cleanup attachment alive), so without the post-unwind + ``task.cancelling()`` re-check the node would end as a failed tool call + instead of an interrupted run.""" + config = _make_subagent_config() + cancel_calls: list[str] = [] + emit_calls = 0 + unwind_entered = asyncio.Event() + execution_id = "exec-grace-cancel" + + async def fail_on_status_emit(event, *, writer=None): + nonlocal emit_calls + emit_calls += 1 + if emit_calls >= 2: + raise RuntimeError("status emit boom") + return None + + async def absorbing_finalize(runtime_arg, execution_id_arg, trace_id_arg, max_polls, grace_seconds=None): + # Mimic the production unwind: park inside the grace wait and absorb + # the outer cancellation, then return so the tool's generic-error + # branch continues past the unwind. + unwind_entered.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + pass + + monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) + monkeypatch.setattr( + task_tool_module, + "SubagentExecutor", + type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: execution_id}), + ) + monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) + monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _: _make_result(FakeSubagentStatus.RUNNING, ai_messages=[])) + monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: lambda _event: None) + monkeypatch.setattr(task_tool_module, "aemit_custom_event", fail_on_status_emit) + monkeypatch.setattr(task_tool_module, "request_cancel_background_task", lambda execution_id_arg: cancel_calls.append(execution_id_arg)) + monkeypatch.setattr(task_tool_module, "_finalize_interrupted_subagent", absorbing_finalize) + monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep) + monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: []) + + coroutine = getattr(task_tool_module.task_tool, "coroutine", None) + assert coroutine is not None + tool_task = asyncio.create_task( + coroutine( + runtime=_make_runtime(), + description="test", + prompt="p", + subagent_type="general-purpose", + tool_call_id="tc-grace-cancel", + ) + ) + await asyncio.wait_for(unwind_entered.wait(), timeout=10.0) + tool_task.cancel() + with pytest.raises(asyncio.CancelledError): + await tool_task + assert cancel_calls == [execution_id] + + +@pytest.mark.asyncio +async def test_deferred_cleanup_does_not_retain_runtime(monkeypatch): + """The deferred cleaner must retain only the resolved usage recorder and + ids — never the whole ``runtime``. The strongly-referenced cleanup task + lives for up to the full poll budget; through ``runtime`` it would pin + the parent run's journal and event store for that entire window, worst + on the polling-timeout path where a stuck subagent pins its run's + journal for a second full timeout after the tool already returned.""" + orig_sleep = asyncio.sleep + current_loop = asyncio.get_running_loop() + + def transport_to_current_loop(coro): + return asyncio.run_coroutine_threadsafe(coro, current_loop) + + monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) + # Never terminal: the cleaner keeps polling for its whole budget. + monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _: _make_result(FakeSubagentStatus.RUNNING, ai_messages=[])) + monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda _: None) + monkeypatch.setattr(task_tool_module, "run_on_isolated_subagent_loop", transport_to_current_loop) + monkeypatch.setattr(task_tool_module.asyncio, "sleep", lambda _: orig_sleep(0)) + + # SimpleNamespace rejects weakref; a plain class instance does not. Only + # runtime.config is read on the scheduling path (recorder resolution). + class WeakrefableRuntime: + config = {"metadata": {"model_name": "ark-model", "trace_id": "trace-1"}} + + runtime = WeakrefableRuntime() + runtime_ref = weakref.ref(runtime) + handle = task_tool_module._schedule_deferred_subagent_cleanup(runtime, "exec-retain", "trace-retain", 50) + assert handle in task_tool_module._deferred_cleanup_tasks + del runtime + gc.collect() + + assert runtime_ref() is None, "deferred cleanup must not pin the run runtime (journal + event store) for its poll budget" + + # Let the cleaner exhaust its budget (no-op sleeps) so the handle settles. + await asyncio.wait_for(asyncio.wrap_future(handle), timeout=10.0) + task_tool_module._deferred_cleanup_tasks.discard(handle) + + +def test_execute_async_failure_leaves_no_background_residue(monkeypatch): + """``execute_async`` raises before the poller's guarded region starts + (e.g. the persistent subagent loop failed to spin up). The registry entry + is rolled back inside ``execute_async`` itself (see #5086), so the tool + must simply propagate the error — no background entry, no deferred + cleanup dependency on the loop that just failed to start.""" + config = _make_subagent_config() + deferred_schedules: list = [] + cleanup_calls: list[str] = [] + monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) + monkeypatch.setattr( + task_tool_module, + "SubagentExecutor", + type( + "FailingExecutor", + (), + { + "__init__": lambda self, **kwargs: None, + "execute_async": lambda self, prompt, task_id=None: (_ for _ in ()).throw(RuntimeError("Timed out starting isolated subagent event loop")), + }, + ), + ) + monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) + monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda execution_id_arg: cleanup_calls.append(execution_id_arg)) + monkeypatch.setattr( + task_tool_module, + "run_on_isolated_subagent_loop", + lambda coro: deferred_schedules.append(coro) or (_ for _ in ()).throw(AssertionError("deferred cleanup must not be scheduled")), + ) + monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: []) + + with pytest.raises(RuntimeError, match="Timed out starting isolated subagent event loop"): + _run_task_tool( + runtime=_make_runtime(), + description="test", + prompt="p", + subagent_type="general-purpose", + tool_call_id="tc-executor-submit-failure", + ) + # Nothing was registered (execute_async rolls its own entry back before + # re-raising), so nothing may be cancelled, cleaned, or deferred either — + # the guarded region's invariants hold vacuously before it starts. + assert deferred_schedules == [] + assert cleanup_calls == [] + + def test_cancelled_cleanup_stops_after_timeout(monkeypatch): """Verify cancellation handler survives a shielded-wait timeout gracefully. @@ -1508,7 +2267,7 @@ def test_cancelled_cleanup_stops_after_timeout(monkeypatch): def add_done_callback(self, callback): self.callback = callback - def fake_create_task(coro): + def fake_run_on_isolated_subagent_loop(coro): scheduled_cleanups.append(coro) coro.close() return DummyCleanupTask(coro) @@ -1522,7 +2281,7 @@ def test_cancelled_cleanup_stops_after_timeout(monkeypatch): monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: events.append) monkeypatch.setattr(task_tool_module.asyncio, "sleep", cancel_on_first_sleep) - monkeypatch.setattr(task_tool_module.asyncio, "create_task", fake_create_task) + monkeypatch.setattr(task_tool_module, "run_on_isolated_subagent_loop", fake_run_on_isolated_subagent_loop) monkeypatch.setattr(task_tool_module, "_report_subagent_usage", fake_report_subagent_usage) monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: []) monkeypatch.setattr( @@ -1966,16 +2725,25 @@ def test_terminal_event_usage_none_when_no_records(monkeypatch): @pytest.mark.asyncio async def test_deferred_cleanup_task_retained_and_survives_gc(monkeypatch): - """Verify deferred cleanup task is retained in _deferred_cleanup_tasks and completes after GC.""" + """Verify deferred cleanup is retained in _deferred_cleanup_tasks and completes after GC.""" cleaned = [] orig_sleep = asyncio.sleep + # Route the production transport onto this test's running loop so the + # scheduled coroutine actually executes; the retention and completion + # behavior under test is the real scheduling wrapper. + current_loop = asyncio.get_running_loop() + + def transport_to_current_loop(coro): + return asyncio.run_coroutine_threadsafe(coro, current_loop) + monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _: _make_result(FakeSubagentStatus.COMPLETED, result="ok")) monkeypatch.setattr(task_tool_module, "cleanup_background_task", cleaned.append) + monkeypatch.setattr(task_tool_module, "run_on_isolated_subagent_loop", transport_to_current_loop) monkeypatch.setattr(task_tool_module.asyncio, "sleep", lambda _: orig_sleep(0)) - task = task_tool_module._schedule_deferred_subagent_cleanup("exec-gc", "trace-gc", 5) + task = task_tool_module._schedule_deferred_subagent_cleanup(_make_runtime(), "exec-gc", "trace-gc", 5) assert task in task_tool_module._deferred_cleanup_tasks weak_task = weakref.ref(task) del task @@ -1989,7 +2757,18 @@ async def test_deferred_cleanup_task_retained_and_survives_gc(monkeypatch): await orig_sleep(0.01) assert cleaned == ["exec-gc"] + # The transport runs on this test's own loop, so the production + # add_done_callback(_deferred_cleanup_tasks.discard) fires on the same + # loop — deterministic once `cleaned` was observed. The assert (not a + # manual discard) is the point: if that done-callback were deleted from + # _schedule_deferred_subagent_cleanup, the handle would linger and this + # would fail instead of the test silently cleaning up after itself. + for _ in range(10): + if weak_task() not in task_tool_module._deferred_cleanup_tasks: + break + await orig_sleep(0.01) assert weak_task() not in task_tool_module._deferred_cleanup_tasks + task_tool_module._deferred_cleanup_tasks.discard(weak_task()) def _receipt_fixture(rid: str = "r1", tool: str = "write_file", status: str = "success") -> dict: