mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-09 21:49:37 +00:00
fix(subagents): clean up background task entry on unexpected poller exit (#5069)
* fix(subagents): clean up background task entry on unexpected poller exit * fix(subagents): pin deferred cleanup to the persistent subagent loop The non-terminal fallback scheduled the deferred registry cleanup with asyncio.create_task on the poller's own loop. Under synchronous tool invocation the sync wrapper runs the tool coroutine through asyncio.run(), which cancels caller-loop tasks at teardown, so the cleanup died before executing and the _background_tasks entry leaked — the same lifecycle leak the terminal path already fixed. Schedule the deferred cleaner on the process-owned persistent subagent loop instead, via the new public executor helper run_on_isolated_subagent_loop (asyncio.run_coroutine_threadsafe). The cleaner only touches thread-safe registry helpers, so it is loop-agnostic. A caller-loop fallback remains for the unreachable case where the persistent loop cannot be obtained, so scheduling never raises out of an unwind path that is already handling an error. The polling-timeout return path, which shares the scheduler, is fixed the same way. Tests no longer stub the scheduler: the non-terminal fallback test drives the real scheduling wrapper on an equivalent long-lived loop and asserts cleanup runs after asyncio.run() tears the caller loop down, and run_on_isolated_subagent_loop itself is covered against the real persistent loop with the caller loop closed underneath. * fix(subagents): harden interrupted finalization against failing status path Three edge cases from review on the unexpected-exit unwind: 1. Finalization no longer depends on the failing status accessor. _peek_subagent_result distinguishes a gone entry from an unreadable one instead of letting the accessor's exception abort the unwind; _finalize_interrupted_subagent never raises (so the original poller exception is preserved) and attaches the deferred cleaner, whose last resort force-removes a persistently unreadable entry via the new executor force_cleanup_background_task. 2. The generic-error unwind waits only a short grace period (_UNEXPECTED_EXIT_GRACE_SECONDS) instead of the full execution timeout before re-raising; the remaining lifecycle stays with the deferred cleaner on the persistent subagent loop. 3. The deferred cleaner reports the subagent's final usage (deltas since the unwind snapshot included, via final=True bypassing usage_reported; the journal dedupes by source_run_id) before removing the terminal entry. The report is transferred in a plain worker thread so the RunJournal's loop-bound progress flush is skipped rather than scheduled on a foreign loop. * fix(subagents): pin deferred final usage delivery to the parent run loop _report_deferred_final_usage ran record_external_llm_usage_records in a to_thread worker, making it the first cross-thread RunJournal writer: the unlocked accumulators can lose token updates and _tokens_by_model mutations race get_completion_data() iteration on the parent loop. Capture the parent loop at unwind time (it is alive in every path that continues the run) and deliver the final report onto it with call_soon_threadsafe, serialized with all other journal access; when that loop is already closed (asyncio.run teardown) the report is dropped on purpose — the run has persisted and nothing reads the counters back. The live-loop test exercises the real recorder path (journal captures the running loop of every call) instead of stubbing _report_subagent_usage, so a cross-thread report would surface as a wrong-loop entry. Also gates the two teardown tests on the caller loop actually closing (the deferred cleaner could otherwise legitimately deliver while that loop is still winding down), and adds a task_tool-level regression test for execute_async submit failure leaving no registry residue (rolled back inside execute_async since #5086). * docs(subagents): document the reverse loop boundary; observability + test fixes Extend subagents/AGENTS.md's Isolated-loop callback boundary with the reverse-direction contract from #5069: deferred registry cleanup is pinned to the persistent subagent loop via run_on_isolated_subagent_loop, and the final usage report is handed back onto the parent run's loop captured at unwind time — never invoked from the persistent loop or a worker thread, which would silently reintroduce the journal accumulator/iteration race. Dropping the final report (closed parent loop) is the one path where a subagent's tail usage goes permanently unaccounted, so both drop branches now log at info with the execution id and the unaccounted record count. Restore the retention assert in test_deferred_cleanup_task_retained_and_ survives_gc: the bounded wait observes the production done-callback discard, the assert (not a manual discard) is what fails if that callback is deleted. * fix(subagents): honour grace-wait cancellation and shrink deferred-cleaner captures Two review follow-ups on the unwind: - The shared unwind absorbs CancelledError (never-raise contract), so a graph-node cancellation landing inside the generic-error grace wait was swallowed and the node ended as a failed tool call instead of an interrupted run. The generic branch now re-checks task.cancelling() after the unwind and re-raises CancelledError; the absorb site documents why the cancellation path needs it and where the discrimination lives. - The deferred cleaner captured the whole run runtime, pinning the parent run's journal and event store for up to a full poll budget (~31 min) via the strongly-held task handle — worst on the polling-timeout path, where a stuck subagent pinned its run's journal for a second full timeout after the tool returned. The recorder is now resolved on the unwind path and is the only capture (plus ids and the report loop); a None recorder skips reporting entirely. Both behaviours are regression-tested (red on the previous head, green after): a cancellation parked inside the grace wait surfaces as CancelledError, and the runtime is collectable while the cleaner still polls. The closed-loop drop test now uses a real recorder via runtime.callbacks so the drop it pins is loop-based, not recorder-absence.
This commit is contained in:
hataa
GitHub
parent
3b601922ff
commit
17cac3420a
@ -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.
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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:
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user