mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-27 13:40:37 +00:00
fix(subagents): bound hard timeout terminalization (#7053)
* fix(subagents): bound hard timeout terminalization Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> * fix(subagents): address timeout lifecycle review Bound durable terminal cleanup waits to 60 seconds while unfinished execution resources stay supervised. Preserve terminal results when cancellation races a child failure, and register deferred registry removal when its fence wait expires. Add five lifecycle regression cases and preload the sandbox prewarm dependency in the executor test fixture. Validate 470 focused tests, 261 strict blocking-I/O tests, and 50 repeated regression executions under asyncio debug. --------- Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
RongJie G
GitHub
parent
b8bb728368
commit
cd28a1cfe9
@ -2400,6 +2400,8 @@ Embedded Python uploads use the same captured content for the saved document and
|
||||
返回游标。目录元数据或调用上下文变化时会明确要求重新枚举；详情见
|
||||
[文件上传文档](backend/docs/FILE_UPLOAD.md)。
|
||||
|
||||
Native subagent hard timeouts publish `TIMED_OUT` at the execution deadline while cleanup continues under supervision. Durable batch workers wait up to 60 seconds for cleanup before recording the published outcome and applying the configured retry policy. Unfinished executions retain capacity and sandbox leases until teardown finishes; keep work idempotent because retries may overlap late cleanup.
|
||||
|
||||
Durable `batch_task` workers use one app-owned plugin snapshot for tool assembly and execution. Recovered tasks adopt the new worker's plugin snapshot after a Gateway restart; plugin objects are never stored in durable task records.
|
||||
|
||||
Ordinary `task` delegation and explicit durable `batch_task` execution share the startup-scoped `subagent_runtime` process capacity. Batch mode keeps large independent item sets in SQL with separate total, live, and running limits, restart recovery, bounded results, and a thread-scoped Web UI panel. The panel pages through bounded previews on demand; bulk stored results are available through the owner-scoped JSONL export, while internal execution and authorization context never enters owner-facing responses. If the batch worker is later stopped or disabled, threads with persisted batches retain read-only item inspection and JSONL export; execution controls remain disabled until the worker is running again. See `config.example.yaml` and [the implementation contract](docs/plans/2026-08-24-subagent-batch-capacity-implementation.md) for limits and recovery semantics.
|
||||
|
||||
@ -76,7 +76,7 @@ threaded initialization, peer cancellation and renewal drain.
|
||||
**Knowledge-scope boundary**: `task` and `batch_task` inherit the parent's canonical execution scope, never display data. Batch specs persist it across lease recovery; malformed scope fails. Model input cannot set or broaden it, and shared middleware/retrieval enforce `disabled`.
|
||||
**Date context (#4781)**: Built-in executions register `SubagentDateContextMiddleware` immediately before `SystemMessageCoalescingMiddleware`. Its one-time `before_agent` injects only a hidden framework `SystemMessage` containing `<current_date>` before the first model call: no `AppConfig.memory`, memory-manager call, task `HumanMessage` rewrite or inherited lead frozen-conversation/midnight lifecycle. Coalescing with the static prompt yields one leading `SystemMessage` for strict providers. Lead-only `DynamicContextMiddleware` retains its date, optional-memory and midnight-update behavior.
|
||||
**Execution**: Ordinary and durable-batch native subagents share one persistent isolated event loop. Retain loop/thread ownership until the worker stops; bounded joins do not prove termination. Fence submissions during cleanup. After exit, the next submission reaps/closes the old loop under the shutdown lifecycle lock before replacement. Gateway/embedded startup installs one process-wide async FIFO admission controller (default 3 running, bounded queue). Direct `create_deerflow_agent` callers may pass a caller-owned `SubagentRuntime`; reuse it across graphs so bound `task`, optional batch tools/service, middleware limits and `SubagentExecutor` share one controller without global YAML. Start owned batch services before graph construction; stop them at application shutdown. Waiters hold no scheduler thread; cancellation/timeout release queue/slot ownership.
|
||||
**Shared sandbox execution lifecycle** (#5128): every admitted subagent run carries a stable task-derived `sandbox_lease_owner_id` and matching `sandbox_command_scope_id` in its runtime context. Sandbox middleware retains that execution against the lead thread's active provider client, so one child finishing cannot close the sandbox while siblings still run; the final holder performs any pending provider release. A rollback/fork-restored child reusing the parent's live client binds a non-releasing holder: it fences parent cleanup and owns its command scope without requesting a park itself; a parent's earlier park request waits for the child, while a missing inherited client falls through to a normal fresh acquire. On AIO, the command scope selects one explicit persistent shell session per subagent, allowing independent scopes to run concurrently while preserving in-order shell state within one child. Sync sandbox tool bodies offloaded with `asyncio.to_thread` are shielded and drained across repeated cancellation before the outer execution can clean its holder; a cancelled worker can therefore neither re-admit an already-released owner nor run after subagent terminalization. `SubagentExecutor` invokes active LangGraph stream cleanup before releasing its sandbox lease or notifying task stop. Slow cooperative cleanup keeps the result non-terminal and retains its sandbox lease and capacity slot; its warning is diagnostic, not a safe hard-timeout boundary. Middleware performs the normal release, and `SubagentExecutor` repeats it idempotently in `finally` so exceptions, cooperative cancellation, and timeout unwind paths cannot leak a lease or scoped session.
|
||||
**Shared sandbox execution lifecycle** (#5128): Admitted subagents carry stable task-derived `sandbox_lease_owner_id` and `sandbox_command_scope_id` in runtime context. Middleware retains each execution on the lead thread's provider client; only the final holder performs pending release. Rollback/fork-restored children bind non-releasing holders on the parent's live client, fence parent parking, and own their command scope; without an inherited client, acquire normally. AIO maps each scope to an explicit persistent shell session: children run concurrently, while commands within one child stay ordered. Shield and drain `asyncio.to_thread` sandbox tools across repeated cancellation before holder cleanup; cancelled workers must neither re-admit released owners nor outlive execution. Close the LangGraph stream before releasing sandbox leases or notifying task stop. Slow cooperative cleanup keeps results non-terminal and retains capacity and sandbox leases; warnings are diagnostic. Hard timeouts publish `TIMED_OUT` immediately; `execution_done_event` records full teardown and gates registry removal. Batch finalization waits at most 60s, then records the published outcome with normal retry policy; idempotent work may overlap late cleanup. Retain resources until teardown actually finishes. Bounded usage-tail cleaners must request registry removal even when their wait expires. Ordinary `task` delivers usage on the parent loop. Middleware releases normally; the executor repeats release idempotently in `finally` for errors, cancellation and timeout.
|
||||
**Concurrency and total delegation cap**: Ordinary `task` concurrency is resolved once as the minimum of the per-run request, the startup-frozen `subagent_runtime.max_running`, and the schema safety ceiling (1-64), then shared by the lead prompt and `SubagentLimitMiddleware`. Hot reloads must not make either layer advertise more capacity than the already-created process controller; a changed startup-only value takes effect only after restart. The same middleware separately enforces `subagents.max_total_per_run` (default 6, config schema 1-50, runtime override `max_total_subagents` clamped to the same range) against current-run entries in the durable delegation ledger, so a long lead-agent run cannot bypass concurrency limits by launching repeated legal-sized batches at each planning checkpoint, but historical delegations from previous runs in the same thread do not consume the new run's budget. Explicit `batch_task` work does not consume or relax that ordinary-run ledger: its persisted total/live/running limits live under `subagent_batches`. Gateway `run_agent()` and embedded `DeerFlowClient.stream()` both provide a per-invocation `run_id` in runtime context; `DeerFlowClient.stream()` also tags its input `HumanMessage` with that same id so durable-context capture can identify the current request boundary. Gateway resume paths may not append a new `HumanMessage`, so the worker also exposes the pre-run checkpoint's message ids in runtime context; durable-context capture uses that as the current-run boundary and never re-tags older task calls as the resumed run. When no delegation slots remain, task calls are stripped, provider raw tool-call metadata is synced, `finish_reason` is forced to `stop`, and a visible "subagent delegation limit" note is appended so the agent can synthesize already-collected results. Default subagent timeout `subagents.timeout_seconds=1800` (30 min) and built-in `general-purpose` `max_turns=150`.
|
||||
**Flow**: `task()` → `SubagentExecutor` → shared slot → polling/SSE. Explicit `batch_task()` → durable rows → lease-based batch service (`subagents/batch_service.py`, Gateway or direct runtime) → same executor/slots → bounded result and owner-scoped API/JSONL. Batch mode is explicit, never inferred from prompt size. Queue rejection/timeout releases the lease without an item attempt; execution failure/expired leases consume one. User cancellation terminalizes nonterminal items and clears leases, fencing stale workers. Resolve background cancellation under `_background_tasks_lock`, then call `Future.cancel()` outside it because callbacks may reacquire the lock. `task_started` includes the effective model. `SubagentTokenCollector` publishes cumulative usage after each LLM response; `task_running` updates cards without double-counting. Terminal ToolMessage metadata (`subagent_model_name`, `subagent_token_usage`) and `subagent.end` preserve model/usage after reload; absent usage stays absent. Resolve one `AppConfig` snapshot (explicit or `get_app_config()`) for assembly, deferred setup, and receipt harvesting; `verification.receipts_enabled=false` holds on both construction paths. Harvest terminal receipts before `try_set_terminal` under the result lock so polling cannot see terminal status without metadata. Each values chunk updates the latest harvest state and publishes receipts before cooperative cancellation. Cancellation/failure scans current ToolMessages; completion uses the bounded ledger snapshot attached to returned assistant text, including max-turn partials. Missing/malformed completed snapshots yield no receipts, preventing omitted receipts from validating citations. Cancellation/timeouts retain the latest streamed evidence.
|
||||
**Report contract (RFC #4651 PR3)**: `report_contract.py` owns the prompt-layer text that makes Layer 1 receipt verification non-inert. `SubagentExecutor._build_initial_state` appends `build_report_contract_section(receipts_enabled=...)` to every subagent's consolidated `SystemMessage` — built-in and custom alike — requiring `[rN tool_name]` citations (from the Tool receipts ledger) for action claims, verifiable handles (absolute path, URL, ID, HTTP status) for deliverables, and explicit reporting of failures; the citation clause follows `verification.receipts_enabled`, and the citation example derives from the single-owner `format_citation`/`receipt_id` so prompt text cannot drift from the verifier. The `task` tool hands lead-supplied `acceptance_criteria` to the `SubagentExecutor` constructor, which appends them via `render_acceptance_criteria_block(...)` to the task `HumanMessage` (stripped, capped at 20 items × 500 chars, each entry neutralized) — the untrusted channel `InputSanitizationMiddleware` escapes and boundary-frames, matching their model-supplied provenance. The subagent's `SystemMessage` never carries criterion text; it gets only the framework-owned `build_acceptance_criteria_system_note(...)` pointer naming the list's location and authority, so natural-language injection inside a criterion cannot gain system-channel priority over framework instructions. Deterministic leaf checking is a separate layer.
|
||||
|
||||
@ -36,6 +36,10 @@ logger = logging.getLogger(__name__)
|
||||
# weak references to its tasks.
|
||||
_pending_releases: set[Future[None]] = set()
|
||||
|
||||
# Terminal outcomes must not keep durable attempts leased forever when resource
|
||||
# teardown hangs. The execution supervisor still owns those resources afterwards.
|
||||
_TERMINAL_CLEANUP_GRACE_SECONDS = 60.0
|
||||
|
||||
|
||||
def _usage(records: list[dict[str, Any]] | None) -> dict[str, int] | None:
|
||||
if not records:
|
||||
@ -361,6 +365,8 @@ class SubagentBatchService:
|
||||
self._execution_ids[item_id] = execution_id
|
||||
marked_running = False
|
||||
status_poll_every = min(self._config.poll_interval_seconds, max(1.0, self._config.lease_seconds / 3))
|
||||
cleanup_deadline: float | None = None
|
||||
loop = asyncio.get_running_loop()
|
||||
while True:
|
||||
result = get_background_task_result(execution_id)
|
||||
if result is None:
|
||||
@ -373,12 +379,27 @@ class SubagentBatchService:
|
||||
)
|
||||
if not marked_running:
|
||||
request_cancel_background_task(execution_id)
|
||||
status_wait = status_poll_every
|
||||
execution_done = getattr(result, "execution_done_event", None)
|
||||
if result.status.is_terminal:
|
||||
break
|
||||
if execution_done is None or execution_done.is_set():
|
||||
break
|
||||
if cleanup_deadline is None:
|
||||
cleanup_deadline = loop.time() + _TERMINAL_CLEANUP_GRACE_SECONDS
|
||||
remaining = cleanup_deadline - loop.time()
|
||||
if remaining <= 0:
|
||||
logger.warning(
|
||||
"Batch execution %s cleanup exceeded %.0fs; finalizing its published %s outcome while resource cleanup remains supervised",
|
||||
execution_id,
|
||||
_TERMINAL_CLEANUP_GRACE_SECONDS,
|
||||
result.status.value,
|
||||
)
|
||||
break
|
||||
status_wait = min(status_wait, remaining)
|
||||
if not lease.valid:
|
||||
request_cancel_background_task(execution_id)
|
||||
try:
|
||||
await asyncio.wait_for(self._stop.wait(), timeout=status_poll_every)
|
||||
await asyncio.wait_for(self._stop.wait(), timeout=status_wait)
|
||||
if self._stop.is_set():
|
||||
raise asyncio.CancelledError
|
||||
except TimeoutError:
|
||||
|
||||
@ -156,6 +156,9 @@ class SubagentResult:
|
||||
execution. ``None`` when the delegation carried no acceptance
|
||||
criteria, the run ended before streaming, or harvesting failed;
|
||||
an empty list means the stream carried no bash-family tool calls.
|
||||
execution_done_event: Internal lifecycle fence set only after stream,
|
||||
sandbox, extension, and capacity cleanup has settled. A configured
|
||||
timeout may publish ``TIMED_OUT`` before this fence opens.
|
||||
"""
|
||||
|
||||
task_id: str
|
||||
@ -176,17 +179,23 @@ class SubagentResult:
|
||||
knowledge_sources: list[dict[str, Any]] = field(default_factory=list, kw_only=True)
|
||||
_knowledge_messages: list[Any] = field(default_factory=list, init=False, repr=False)
|
||||
cancel_event: threading.Event = field(default_factory=threading.Event, repr=False)
|
||||
execution_done_event: threading.Event = field(default_factory=threading.Event, init=False, repr=False)
|
||||
_state_lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
|
||||
|
||||
def __post_init__(self):
|
||||
"""Initialize mutable defaults."""
|
||||
if self.ai_messages is None:
|
||||
self.ai_messages = []
|
||||
if self.status.is_terminal or self.completed_at is not None:
|
||||
self.execution_done_event.set()
|
||||
|
||||
def update_token_usage_records(self, records: list[dict[str, int | str | None]]) -> None:
|
||||
"""Publish the latest cumulative collector snapshot while still running."""
|
||||
"""Publish cumulative usage, including timeout/cancel cleanup tails."""
|
||||
with self._state_lock:
|
||||
if not self.status.is_terminal:
|
||||
if not self.status.is_terminal or self.status in {
|
||||
SubagentStatus.CANCELLED,
|
||||
SubagentStatus.TIMED_OUT,
|
||||
}:
|
||||
self.token_usage_records = list(records)
|
||||
|
||||
def update_tool_receipts(self, receipts: list[dict[str, Any]] | None) -> None:
|
||||
@ -411,6 +420,10 @@ _background_tasks: dict[str, SubagentResult] = {}
|
||||
_background_tasks_lock = threading.Lock()
|
||||
|
||||
_background_futures: dict[str, Future[SubagentResult]] = {}
|
||||
_background_cleanup_requested: set[str] = set()
|
||||
# A timed-out public result can precede stream/resource cleanup. Keep the
|
||||
# execution task alive and observed until that cleanup settles on its owner loop.
|
||||
_background_execution_tasks: set[asyncio.Task[SubagentResult]] = set()
|
||||
|
||||
|
||||
def _harvest_tool_receipts(
|
||||
@ -1620,9 +1633,10 @@ class SubagentExecutor:
|
||||
capacity = self.execution_capacity or get_subagent_execution_capacity()
|
||||
async with capacity.slot():
|
||||
with result._state_lock:
|
||||
if not result.status.is_terminal:
|
||||
result.status = SubagentStatus.RUNNING
|
||||
result.started_at = _utcnow()
|
||||
if result.status.is_terminal:
|
||||
return result
|
||||
result.status = SubagentStatus.RUNNING
|
||||
result.started_at = _utcnow()
|
||||
return await self._aexecute_admitted(task, result)
|
||||
except SubagentCapacityError as exc:
|
||||
result.try_set_terminal(
|
||||
@ -1631,6 +1645,8 @@ class SubagentExecutor:
|
||||
admission_failure=True,
|
||||
)
|
||||
return result
|
||||
finally:
|
||||
result.execution_done_event.set()
|
||||
|
||||
async def _aexecute_admitted(self, task: str, result_holder: SubagentResult | None = None) -> SubagentResult:
|
||||
"""Execute a task asynchronously.
|
||||
@ -2200,21 +2216,46 @@ class SubagentExecutor:
|
||||
_background_tasks[execution_id] = result
|
||||
|
||||
async def run_with_timeout() -> SubagentResult:
|
||||
execution_task = asyncio.create_task(
|
||||
self._aexecute(task, result),
|
||||
name=f"subagent-{self.config.name}-execution",
|
||||
)
|
||||
try:
|
||||
return await asyncio.wait_for(
|
||||
self._aexecute(task, result),
|
||||
done, _ = await asyncio.wait(
|
||||
{execution_task},
|
||||
timeout=self.config.timeout_seconds,
|
||||
)
|
||||
except TimeoutError:
|
||||
if execution_task in done:
|
||||
return execution_task.result()
|
||||
|
||||
# Publish the configured hard timeout without waiting for
|
||||
# cancellation-resistant stream cleanup. The detached execution
|
||||
# still owns its capacity slot and sandbox lease until its
|
||||
# supervisor observes full teardown.
|
||||
result.cancel_event.set()
|
||||
result.try_set_terminal(
|
||||
SubagentStatus.TIMED_OUT,
|
||||
error=f"Execution timed out after {self.config.timeout_seconds} seconds",
|
||||
tool_receipts=result.snapshot_tool_receipts(),
|
||||
)
|
||||
execution_task.cancel()
|
||||
_supervise_background_execution(
|
||||
execution_id,
|
||||
result,
|
||||
execution_task,
|
||||
)
|
||||
execution_task = None
|
||||
return result
|
||||
except asyncio.CancelledError:
|
||||
result.cancel_event.set()
|
||||
if not execution_task.done():
|
||||
execution_task.cancel()
|
||||
try:
|
||||
await execution_task
|
||||
except (asyncio.CancelledError, Exception):
|
||||
# The child can fail before cancellation reaches it. Keep
|
||||
# its terminal write, or publish cancellation below.
|
||||
pass
|
||||
result.try_set_terminal(
|
||||
SubagentStatus.CANCELLED,
|
||||
error="Cancelled by user",
|
||||
@ -2223,8 +2264,19 @@ class SubagentExecutor:
|
||||
return result
|
||||
except Exception as exc:
|
||||
logger.exception("[trace=%s] Subagent %s async execution failed", self.trace_id, self.config.name)
|
||||
if not execution_task.done():
|
||||
execution_task.cancel()
|
||||
_supervise_background_execution(
|
||||
execution_id,
|
||||
result,
|
||||
execution_task,
|
||||
)
|
||||
execution_task = None
|
||||
result.try_set_terminal(SubagentStatus.FAILED, error=str(exc))
|
||||
return result
|
||||
finally:
|
||||
if execution_task is not None:
|
||||
_mark_background_execution_done(execution_id, result)
|
||||
|
||||
try:
|
||||
execution_future = _submit_to_isolated_loop_in_context(parent_context, run_with_timeout)
|
||||
@ -2241,7 +2293,9 @@ class SubagentExecutor:
|
||||
with _background_tasks_lock:
|
||||
_background_futures[execution_id] = execution_future
|
||||
|
||||
def forget_future(_future: Future[SubagentResult]) -> None:
|
||||
def forget_future(completed_future: Future[SubagentResult]) -> None:
|
||||
if not completed_future.cancelled():
|
||||
completed_future.exception()
|
||||
with _background_tasks_lock:
|
||||
_background_futures.pop(execution_id, None)
|
||||
|
||||
@ -2298,14 +2352,55 @@ def list_background_tasks() -> list[SubagentResult]:
|
||||
return list(_background_tasks.values())
|
||||
|
||||
|
||||
def _remove_background_task_locked(execution_id: str) -> None:
|
||||
_background_tasks.pop(execution_id, None)
|
||||
_background_futures.pop(execution_id, None)
|
||||
_background_cleanup_requested.discard(execution_id)
|
||||
|
||||
|
||||
def _mark_background_execution_done(
|
||||
execution_id: str,
|
||||
result: SubagentResult,
|
||||
) -> None:
|
||||
result.execution_done_event.set()
|
||||
with _background_tasks_lock:
|
||||
if execution_id in _background_cleanup_requested and _background_tasks.get(execution_id) is result:
|
||||
_remove_background_task_locked(execution_id)
|
||||
|
||||
|
||||
def _supervise_background_execution(
|
||||
execution_id: str,
|
||||
result: SubagentResult,
|
||||
task: asyncio.Task[SubagentResult],
|
||||
) -> None:
|
||||
"""Own a timed-out/cancelled execution until resource teardown settles."""
|
||||
_background_execution_tasks.add(task)
|
||||
|
||||
def settled(completed: asyncio.Task[SubagentResult]) -> None:
|
||||
_background_execution_tasks.discard(completed)
|
||||
if not completed.cancelled():
|
||||
try:
|
||||
completed.result()
|
||||
except BaseException as exc:
|
||||
logger.error(
|
||||
"Detached cleanup for background execution %s failed",
|
||||
execution_id,
|
||||
exc_info=(type(exc), exc, exc.__traceback__),
|
||||
)
|
||||
_mark_background_execution_done(execution_id, result)
|
||||
|
||||
task.add_done_callback(settled)
|
||||
|
||||
|
||||
def cleanup_background_task(execution_id: str) -> None:
|
||||
"""Remove a completed task from background tasks.
|
||||
|
||||
Should be called by task_tool after it finishes polling and returns the result.
|
||||
This prevents memory leaks from accumulated completed tasks.
|
||||
|
||||
Only removes tasks that are in a terminal state (COMPLETED/FAILED/TIMED_OUT)
|
||||
to avoid race conditions with the background executor still updating the task entry.
|
||||
Only removes a terminal task after its execution-cleanup fence opens. A hard
|
||||
timeout may publish its business outcome first; cleanup requested in that
|
||||
window is completed by the detached execution supervisor.
|
||||
|
||||
Args:
|
||||
execution_id: The execution ID to remove.
|
||||
@ -2314,15 +2409,20 @@ def cleanup_background_task(execution_id: str) -> None:
|
||||
result = _background_tasks.get(execution_id)
|
||||
if result is None:
|
||||
# Nothing to clean up; may have been removed already.
|
||||
_background_cleanup_requested.discard(execution_id)
|
||||
logger.debug("Requested cleanup for unknown background execution %s", execution_id)
|
||||
return
|
||||
|
||||
# Only clean up tasks that are in a terminal state to avoid races with
|
||||
# the background executor still updating the task entry.
|
||||
if result.status.is_terminal or result.completed_at is not None:
|
||||
del _background_tasks[execution_id]
|
||||
_background_futures.pop(execution_id, None)
|
||||
terminal = result.status.is_terminal or result.completed_at is not None
|
||||
if terminal and result.execution_done_event.is_set():
|
||||
_remove_background_task_locked(execution_id)
|
||||
logger.debug("Cleaned up background execution: %s", execution_id)
|
||||
elif terminal:
|
||||
_background_cleanup_requested.add(execution_id)
|
||||
logger.debug(
|
||||
"Deferred cleanup for background execution %s until teardown completes",
|
||||
execution_id,
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
"Skipping cleanup for non-terminal background execution %s (status=%s)",
|
||||
@ -2347,6 +2447,5 @@ def force_cleanup_background_task(execution_id: str) -> None:
|
||||
execution_id: The execution ID to remove.
|
||||
"""
|
||||
with _background_tasks_lock:
|
||||
_background_tasks.pop(execution_id, None)
|
||||
_background_futures.pop(execution_id, None)
|
||||
_remove_background_task_locked(execution_id)
|
||||
logger.warning("Force-cleaned background execution %s after unreadable status", execution_id)
|
||||
|
||||
@ -73,6 +73,11 @@ _SUBAGENT_POLL_INTERVAL_SECONDS = 5.0
|
||||
# cleaner on the persistent subagent loop.
|
||||
_UNEXPECTED_EXIT_GRACE_SECONDS = 5.0
|
||||
|
||||
# A hard timeout returns to the lead immediately while cleanup stays owned on
|
||||
# the isolated loop. Retain its terminal result for this bounded usage-tail
|
||||
# window so callbacks emitted during cleanup are not lost.
|
||||
_TERMINAL_USAGE_TAIL_POLLS = 12
|
||||
|
||||
# Sentinel returned by ``_peek_subagent_result`` when the registry entry exists
|
||||
# but cannot be read (persistent status-lookup / status-object failure).
|
||||
_STATUS_UNREADABLE = object()
|
||||
@ -160,7 +165,20 @@ class _ParentLoopMiddlewareRecorderProxy:
|
||||
|
||||
def _is_subagent_terminal(result: Any) -> bool:
|
||||
"""Return whether a background subagent result is safe to clean up."""
|
||||
return result.status in {SubagentStatus.COMPLETED, SubagentStatus.FAILED, SubagentStatus.CANCELLED, SubagentStatus.TIMED_OUT} or getattr(result, "completed_at", None) is not None
|
||||
status = result.status
|
||||
terminal = getattr(status, "is_terminal", None)
|
||||
if isinstance(terminal, bool):
|
||||
return terminal
|
||||
return (
|
||||
getattr(status, "value", status)
|
||||
in {
|
||||
"completed",
|
||||
"failed",
|
||||
"cancelled",
|
||||
"timed_out",
|
||||
}
|
||||
or getattr(result, "completed_at", None) is not None
|
||||
)
|
||||
|
||||
|
||||
def _peek_subagent_result(execution_id: str, *, trace_id: str) -> Any:
|
||||
@ -290,6 +308,13 @@ async def _finalize_interrupted_subagent(
|
||||
if final_result is not None:
|
||||
_report_subagent_usage(runtime, final_result)
|
||||
if final_terminal:
|
||||
_supervise_terminal_usage_tail(
|
||||
runtime,
|
||||
execution_id,
|
||||
trace_id,
|
||||
max_polls,
|
||||
final_result,
|
||||
)
|
||||
cleanup_background_task(execution_id)
|
||||
else:
|
||||
_schedule_deferred_subagent_cleanup(runtime, execution_id, trace_id, max_polls)
|
||||
@ -372,14 +397,15 @@ async def _deferred_cleanup_subagent_task(
|
||||
max_polls: int,
|
||||
*,
|
||||
report_loop: asyncio.AbstractEventLoop | None = None,
|
||||
terminal_result: Any | 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.
|
||||
Only the resolved usage recorder is retained (plus ids, the captured report
|
||||
loop, and an optional already-terminal result) — 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
|
||||
@ -392,7 +418,9 @@ async def _deferred_cleanup_subagent_task(
|
||||
"""
|
||||
cleanup_poll_count = 0
|
||||
while True:
|
||||
result = _peek_subagent_result(execution_id, trace_id=trace_id)
|
||||
result = terminal_result
|
||||
if result is None:
|
||||
result = _peek_subagent_result(execution_id, trace_id=trace_id)
|
||||
if result is None:
|
||||
return
|
||||
if result is _STATUS_UNREADABLE:
|
||||
@ -401,6 +429,17 @@ async def _deferred_cleanup_subagent_task(
|
||||
force_cleanup_background_task(execution_id)
|
||||
return
|
||||
elif _is_subagent_terminal(result):
|
||||
execution_done = getattr(result, "execution_done_event", None)
|
||||
if isinstance(execution_done, threading.Event) and not execution_done.is_set():
|
||||
if cleanup_poll_count >= max_polls:
|
||||
# Non-pinned callers have not requested registry removal
|
||||
# yet. The supervisor must reap it when the fence opens.
|
||||
cleanup_background_task(execution_id)
|
||||
logger.warning(f"[trace={trace_id}] Deferred cleanup for terminal execution {execution_id} timed out after {cleanup_poll_count} polls")
|
||||
return
|
||||
await asyncio.sleep(_SUBAGENT_POLL_INTERVAL_SECONDS)
|
||||
cleanup_poll_count += 1
|
||||
continue
|
||||
# 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)
|
||||
@ -472,6 +511,8 @@ def _schedule_deferred_subagent_cleanup(
|
||||
execution_id: str,
|
||||
trace_id: str,
|
||||
max_polls: int,
|
||||
*,
|
||||
terminal_result: Any | None = None,
|
||||
) -> asyncio.Task[None] | concurrent.futures.Future:
|
||||
"""Schedule the deferred registry cleanup on the process-owned subagent loop.
|
||||
|
||||
@ -482,7 +523,7 @@ def _schedule_deferred_subagent_cleanup(
|
||||
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}")
|
||||
logger.debug(f"[trace={trace_id}] Scheduling deferred cleanup for execution {execution_id}")
|
||||
# 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
|
||||
@ -494,7 +535,14 @@ def _schedule_deferred_subagent_cleanup(
|
||||
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)
|
||||
coro = _deferred_cleanup_subagent_task(
|
||||
usage_recorder,
|
||||
execution_id,
|
||||
trace_id,
|
||||
max_polls,
|
||||
report_loop=report_loop,
|
||||
terminal_result=terminal_result,
|
||||
)
|
||||
try:
|
||||
cleanup_handle = run_on_isolated_subagent_loop(coro)
|
||||
except Exception:
|
||||
@ -599,6 +647,29 @@ def _report_subagent_usage(runtime: Any, result: Any, *, final: bool = False) ->
|
||||
_report_usage_records(_find_usage_recorder(runtime), result, final=final)
|
||||
|
||||
|
||||
def _supervise_terminal_usage_tail(
|
||||
runtime: Any,
|
||||
execution_id: str,
|
||||
trace_id: str,
|
||||
max_polls: int,
|
||||
result: Any,
|
||||
) -> None:
|
||||
"""Deliver usage collected after a terminal outcome was published."""
|
||||
execution_done = getattr(result, "execution_done_event", None)
|
||||
if not isinstance(execution_done, threading.Event):
|
||||
return
|
||||
if execution_done.is_set():
|
||||
_report_subagent_usage(runtime, result, final=True)
|
||||
return
|
||||
_schedule_deferred_subagent_cleanup(
|
||||
runtime,
|
||||
execution_id,
|
||||
trace_id,
|
||||
max_polls,
|
||||
terminal_result=result,
|
||||
)
|
||||
|
||||
|
||||
def _get_runtime_app_config(runtime: Any) -> "AppConfig | None":
|
||||
explicit = _explicit_app_config.get()
|
||||
if explicit is not None:
|
||||
@ -1205,6 +1276,13 @@ async def task_tool(
|
||||
writer=writer,
|
||||
)
|
||||
logger.warning(f"[trace={trace_id}] Task {tool_call_id} timed out: {result.error}")
|
||||
_supervise_terminal_usage_tail(
|
||||
runtime,
|
||||
execution_id,
|
||||
trace_id,
|
||||
_TERMINAL_USAGE_TAIL_POLLS,
|
||||
result,
|
||||
)
|
||||
cleanup_background_task(execution_id)
|
||||
return _task_result_command(
|
||||
tool_call_id=tool_call_id,
|
||||
|
||||
@ -23,10 +23,11 @@ class FakeStatus(Enum):
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
TIMED_OUT = "timed_out"
|
||||
|
||||
@property
|
||||
def is_terminal(self) -> bool:
|
||||
return self in {FakeStatus.COMPLETED, FakeStatus.FAILED, FakeStatus.CANCELLED}
|
||||
return self in {FakeStatus.COMPLETED, FakeStatus.FAILED, FakeStatus.CANCELLED, FakeStatus.TIMED_OUT}
|
||||
|
||||
|
||||
def _request(**overrides) -> BatchSubmitRequest:
|
||||
@ -401,6 +402,42 @@ async def _wait_until(predicate, *, timeout: float = 2.0) -> None:
|
||||
await asyncio.wait_for(poll(), timeout=timeout)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_terminal_item_waits_for_execution_cleanup_before_finalizing(monkeypatch) -> None:
|
||||
registry = _Registry()
|
||||
terminal_published = asyncio.Event()
|
||||
|
||||
async def mark_item_running():
|
||||
result = registry.results["execution-1"]
|
||||
result.status = FakeStatus.COMPLETED
|
||||
result.result = "done"
|
||||
terminal_published.set()
|
||||
return True
|
||||
|
||||
service, finalized = _running_item_service(
|
||||
monkeypatch,
|
||||
registry,
|
||||
mark_item_running=mark_item_running,
|
||||
)
|
||||
execution_done = asyncio.Event()
|
||||
registry.results["execution-1"].execution_done_event = execution_done
|
||||
|
||||
await service.run_once(now=service_module.datetime.now(service_module.UTC))
|
||||
await asyncio.wait_for(terminal_published.wait(), timeout=1)
|
||||
await asyncio.sleep(0.15)
|
||||
assert finalized.kwargs is None
|
||||
|
||||
execution_done.set()
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(*list(service._executions.values())),
|
||||
timeout=1,
|
||||
)
|
||||
|
||||
assert finalized.kwargs is not None
|
||||
assert finalized.kwargs["succeeded"] is True
|
||||
assert finalized.kwargs["result"] == "done"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stopped_item_releases_its_execution_once_it_turns_terminal(monkeypatch) -> None:
|
||||
# stop() cancels the item task while its subagent is still running, so the
|
||||
@ -977,3 +1014,40 @@ async def test_release_logs_a_failure_on_the_subagent_loop(monkeypatch, caplog)
|
||||
failures = [record for record in caplog.records if "execution-1" in record.getMessage()]
|
||||
assert len(failures) == 1 and failures[0].levelno == logging.ERROR
|
||||
assert str(failures[0].exc_info[1]) == "registry lock poisoned"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_terminal_item_cleanup_fence_budget_does_not_renew_forever(monkeypatch, caplog) -> None:
|
||||
registry = _Registry()
|
||||
execution_done = asyncio.Event()
|
||||
cleanup_requests = []
|
||||
|
||||
async def mark_item_running():
|
||||
result = registry.results["execution-1"]
|
||||
result.status = FakeStatus.TIMED_OUT
|
||||
result.error = "Execution timed out"
|
||||
return True
|
||||
|
||||
service, finalized = _running_item_service(monkeypatch, registry, mark_item_running=mark_item_running)
|
||||
registry.results["execution-1"].execution_done_event = execution_done
|
||||
monkeypatch.setattr(service_module, "_TERMINAL_CLEANUP_GRACE_SECONDS", 0.0, raising=False)
|
||||
# Like the real registry, retain the result while resource teardown is live.
|
||||
monkeypatch.setattr(service_module, "cleanup_background_task", cleanup_requests.append)
|
||||
await service.run_once(now=service_module.datetime.now(service_module.UTC))
|
||||
workers = list(service._executions.values())
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING, logger=service_module.__name__):
|
||||
await asyncio.wait_for(asyncio.gather(*workers), timeout=0.5)
|
||||
assert finalized.kwargs is not None
|
||||
assert finalized.kwargs["succeeded"] is False
|
||||
assert finalized.kwargs["error"] == "Execution timed out"
|
||||
assert not execution_done.is_set()
|
||||
assert registry.get("execution-1") is not None
|
||||
assert cleanup_requests == ["execution-1"]
|
||||
assert registry.forced == []
|
||||
assert "cleanup" in caplog.text.lower()
|
||||
assert not service._execution_ids
|
||||
finally:
|
||||
execution_done.set()
|
||||
await asyncio.gather(*workers, return_exceptions=True)
|
||||
registry.cleanup("execution-1")
|
||||
|
||||
@ -155,6 +155,7 @@ def _setup_executor_classes():
|
||||
original_tool_search = sys.modules.get("deerflow.tools.builtins.tool_search")
|
||||
original_sandbox_provider = sys.modules.get("deerflow.sandbox.sandbox_provider")
|
||||
original_sandbox_overwrite = sys.modules.get("deerflow.sandbox.overwrite")
|
||||
original_sandbox_prewarm = sys.modules.get("deerflow.sandbox.prewarm")
|
||||
|
||||
# Preload real executor dependencies before replacing their parent packages
|
||||
# with cycle-breaking test doubles. Keeping the concrete leaf modules in
|
||||
@ -164,6 +165,7 @@ def _setup_executor_classes():
|
||||
tool_search_module = importlib.import_module("deerflow.tools.builtins.tool_search")
|
||||
sandbox_provider_module = importlib.import_module("deerflow.sandbox.sandbox_provider")
|
||||
sandbox_overwrite_module = importlib.import_module("deerflow.sandbox.overwrite")
|
||||
sandbox_prewarm_module = importlib.import_module("deerflow.sandbox.prewarm")
|
||||
|
||||
# Remove mocked executor if exists (from conftest.py)
|
||||
if "deerflow.subagents.executor" in sys.modules:
|
||||
@ -182,6 +184,7 @@ def _setup_executor_classes():
|
||||
sys.modules["deerflow.tools.builtins.tool_search"] = tool_search_module
|
||||
sys.modules["deerflow.sandbox.sandbox_provider"] = sandbox_provider_module
|
||||
sys.modules["deerflow.sandbox.overwrite"] = sandbox_overwrite_module
|
||||
sys.modules["deerflow.sandbox.prewarm"] = sandbox_prewarm_module
|
||||
|
||||
# Import real classes inside fixture
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
||||
@ -246,6 +249,10 @@ def _setup_executor_classes():
|
||||
sys.modules["deerflow.sandbox.overwrite"] = original_sandbox_overwrite
|
||||
else:
|
||||
sys.modules.pop("deerflow.sandbox.overwrite", None)
|
||||
if original_sandbox_prewarm is not None:
|
||||
sys.modules["deerflow.sandbox.prewarm"] = original_sandbox_prewarm
|
||||
else:
|
||||
sys.modules.pop("deerflow.sandbox.prewarm", None)
|
||||
|
||||
|
||||
def _async_create_agent_double(agent):
|
||||
@ -1611,6 +1618,55 @@ class TestAgentConstruction:
|
||||
class TestAsyncExecutionPath:
|
||||
"""Test _aexecute() async execution path."""
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aexecute_does_not_admit_an_already_terminal_result(
|
||||
self,
|
||||
classes,
|
||||
base_config,
|
||||
):
|
||||
SubagentExecutor = classes["SubagentExecutor"]
|
||||
SubagentResult = classes["SubagentResult"]
|
||||
SubagentStatus = classes["SubagentStatus"]
|
||||
|
||||
class ImmediateCapacity:
|
||||
def slot(self):
|
||||
return self
|
||||
|
||||
async def __aenter__(self):
|
||||
return None
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
result = SubagentResult(
|
||||
task_id="timed-out-before-admission",
|
||||
trace_id="trace",
|
||||
status=SubagentStatus.RUNNING,
|
||||
)
|
||||
assert result.try_set_terminal(
|
||||
SubagentStatus.TIMED_OUT,
|
||||
error="Execution timed out before admission settled",
|
||||
)
|
||||
assert not result.execution_done_event.is_set()
|
||||
|
||||
executor = SubagentExecutor(
|
||||
config=base_config,
|
||||
tools=[],
|
||||
thread_id="test-thread",
|
||||
execution_capacity=ImmediateCapacity(),
|
||||
)
|
||||
admitted = AsyncMock(return_value=result)
|
||||
with patch.object(
|
||||
executor,
|
||||
"_aexecute_admitted",
|
||||
admitted,
|
||||
):
|
||||
returned = await executor._aexecute("Task", result_holder=result)
|
||||
|
||||
assert returned is result
|
||||
admitted.assert_not_awaited()
|
||||
assert result.execution_done_event.is_set()
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aexecute_success(self, classes, base_config, mock_agent, msg):
|
||||
"""Test successful async execution returns completed result."""
|
||||
@ -5738,6 +5794,8 @@ class TestInterruptedTokenUsage:
|
||||
gate["loop"].call_soon_threadsafe(gate["release"].set)
|
||||
assert terminal.wait(10), "subagent did not reach a terminal status"
|
||||
assert result.status.value == outcome
|
||||
if outcome == "timed_out":
|
||||
assert result.execution_done_event.wait(10), "timed-out execution cleanup did not finish"
|
||||
records = result.token_usage_records
|
||||
assert len(records) == 2
|
||||
assert len({record["source_run_id"] for record in records}) == 2
|
||||
@ -6198,6 +6256,113 @@ class TestToolReceiptHarvest:
|
||||
assert result.status == SubagentStatus.TIMED_OUT
|
||||
assert result.tool_receipts == harvested
|
||||
|
||||
def test_execute_async_timeout_terminalizes_before_slow_stream_cleanup(self, classes, monkeypatch):
|
||||
executor_module = importlib.import_module("deerflow.subagents.executor")
|
||||
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
||||
from deerflow.subagents.capacity import SubagentExecutionCapacity
|
||||
|
||||
SubagentExecutor = classes["SubagentExecutor"]
|
||||
SubagentStatus = classes["SubagentStatus"]
|
||||
short_config = classes["SubagentConfig"](
|
||||
name="test-agent",
|
||||
description="Test agent",
|
||||
system_prompt="You are a test agent.",
|
||||
max_turns=10,
|
||||
timeout_seconds=0.05,
|
||||
)
|
||||
stream_started = threading.Event()
|
||||
close_started = threading.Event()
|
||||
release_close = threading.Event()
|
||||
|
||||
class Stream:
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
stream_started.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def aclose(self):
|
||||
close_started.set()
|
||||
while not release_close.is_set():
|
||||
await asyncio.sleep(0.001)
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.astream.return_value = Stream()
|
||||
capacity = SubagentExecutionCapacity(
|
||||
SubagentRuntimeConfig(
|
||||
max_running=1,
|
||||
max_queued=1,
|
||||
queue_timeout_seconds=1,
|
||||
)
|
||||
)
|
||||
executor = SubagentExecutor(
|
||||
config=short_config,
|
||||
tools=[],
|
||||
thread_id="test-thread",
|
||||
extensions=SimpleNamespace(
|
||||
needs_task_store=False,
|
||||
has_task_lifecycle=False,
|
||||
),
|
||||
execution_capacity=capacity,
|
||||
)
|
||||
submitter = self._ThreadSubmitter()
|
||||
task_id = None
|
||||
result = None
|
||||
try:
|
||||
with (
|
||||
patch.object(
|
||||
executor_module,
|
||||
"_submit_to_isolated_loop_in_context",
|
||||
side_effect=submitter.submit,
|
||||
),
|
||||
patch.object(executor_module, "build_tracing_callbacks", return_value=[]),
|
||||
patch.object(executor_module, "inject_langfuse_metadata"),
|
||||
patch.object(
|
||||
executor,
|
||||
"_build_initial_state",
|
||||
new=AsyncMock(return_value=({}, [], None)),
|
||||
),
|
||||
patch.object(
|
||||
executor,
|
||||
"_create_agent",
|
||||
return_value=mock_agent,
|
||||
),
|
||||
):
|
||||
task_id = executor.execute_async("Do something")
|
||||
assert stream_started.wait(timeout=3)
|
||||
assert close_started.wait(timeout=3)
|
||||
|
||||
deadline = time.monotonic() + 1
|
||||
result = executor_module._background_tasks[task_id]
|
||||
while not result.status.is_terminal and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
|
||||
assert result.status is SubagentStatus.TIMED_OUT
|
||||
assert result.cancel_event.is_set()
|
||||
assert not result.execution_done_event.is_set()
|
||||
assert capacity.snapshot().running == 1
|
||||
assert len(executor_module._background_execution_tasks) == 1
|
||||
executor_module.cleanup_background_task(task_id)
|
||||
assert executor_module.get_background_task_result(task_id) is result
|
||||
finally:
|
||||
release_close.set()
|
||||
if result is not None:
|
||||
execution_done = getattr(result, "execution_done_event", None)
|
||||
if execution_done is not None:
|
||||
assert execution_done.wait(timeout=3)
|
||||
elif task_id is not None:
|
||||
deadline = time.monotonic() + 3
|
||||
while task_id in executor_module._background_futures and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
if task_id is not None:
|
||||
executor_module.cleanup_background_task(task_id)
|
||||
submitter.close()
|
||||
|
||||
assert capacity.snapshot().running == 0
|
||||
assert executor_module._background_execution_tasks == set()
|
||||
assert executor_module.get_background_task_result(task_id) is None
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_harvest_skipped_when_receipts_disabled(self, classes, base_config, mock_agent, msg, monkeypatch):
|
||||
SubagentExecutor = classes["SubagentExecutor"]
|
||||
@ -6766,3 +6931,83 @@ async def test_subagent_model_failure_drains_unclaimed_prewarm(classes, base_con
|
||||
assert get_sandbox_lease_manager(provider).binding_for(captured[0]["sandbox_lease_owner_id"]) is None
|
||||
finally:
|
||||
discard_sandbox_lease_manager(provider)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("published_failure", [False, True])
|
||||
async def test_execute_async_cancellation_unwind_child_failure_leaves_terminal_result(classes, base_config, monkeypatch, published_failure):
|
||||
from concurrent.futures import Future
|
||||
|
||||
executor_module = importlib.import_module("deerflow.subagents.executor")
|
||||
executor = classes["SubagentExecutor"](config=base_config, tools=[], thread_id="thread-cancel-race")
|
||||
started = asyncio.Event()
|
||||
fail = asyncio.Event()
|
||||
failed = asyncio.Event()
|
||||
submitted = []
|
||||
future = Future()
|
||||
|
||||
def submit(_context, factory):
|
||||
submitted.append(asyncio.create_task(factory()))
|
||||
return future
|
||||
|
||||
async def execute(_prompt, result):
|
||||
result.status = classes["SubagentStatus"].RUNNING
|
||||
started.set()
|
||||
await fail.wait()
|
||||
if published_failure:
|
||||
result.try_set_terminal(classes["SubagentStatus"].FAILED, error="execution failed")
|
||||
failed.set()
|
||||
raise RuntimeError("execution unwind failed")
|
||||
|
||||
monkeypatch.setattr(executor_module, "_submit_to_isolated_loop_in_context", submit)
|
||||
monkeypatch.setattr(executor, "_aexecute", execute)
|
||||
execution_id = executor.execute_async("work")
|
||||
wrapper = submitted[0]
|
||||
try:
|
||||
async with asyncio.timeout(1):
|
||||
await started.wait()
|
||||
fail.set()
|
||||
# This waiter is queued before asyncio.wait's child-done callback,
|
||||
# so cancel after the child fails but before the wrapper resumes.
|
||||
await failed.wait()
|
||||
wrapper.cancel()
|
||||
result = await wrapper
|
||||
expected = classes["SubagentStatus"].FAILED if published_failure else classes["SubagentStatus"].CANCELLED
|
||||
assert result.status is expected
|
||||
assert result.execution_done_event.is_set()
|
||||
if published_failure:
|
||||
assert result.error == "execution failed"
|
||||
future.set_result(result)
|
||||
executor_module.cleanup_background_task(execution_id)
|
||||
assert executor_module.get_background_task_result(execution_id) is None
|
||||
finally:
|
||||
fail.set()
|
||||
if not wrapper.done():
|
||||
wrapper.cancel()
|
||||
await asyncio.gather(wrapper, return_exceptions=True)
|
||||
if not future.done():
|
||||
future.set_result(None)
|
||||
executor_module.force_cleanup_background_task(execution_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("pinned", [False, True])
|
||||
async def test_deferred_cleanup_expiry_registers_removal_after_fence_opens(classes, monkeypatch, pinned):
|
||||
executor_module = importlib.import_module("deerflow.subagents.executor")
|
||||
tool_module = importlib.import_module("deerflow.tools.builtins.task_tool")
|
||||
execution_id = "execution-expired-cleanup"
|
||||
result = classes["SubagentResult"](task_id=execution_id, trace_id="trace", status=classes["SubagentStatus"].RUNNING)
|
||||
result.try_set_terminal(classes["SubagentStatus"].TIMED_OUT, error="timeout")
|
||||
monkeypatch.setitem(executor_module._background_tasks, execution_id, result)
|
||||
monkeypatch.setattr(tool_module, "SubagentStatus", classes["SubagentStatus"])
|
||||
monkeypatch.setattr(tool_module, "get_background_task_result", executor_module.get_background_task_result)
|
||||
monkeypatch.setattr(tool_module, "cleanup_background_task", executor_module.cleanup_background_task)
|
||||
try:
|
||||
await tool_module._deferred_cleanup_subagent_task(None, execution_id, "trace", 0, terminal_result=result if pinned else None)
|
||||
assert execution_id in executor_module._background_cleanup_requested
|
||||
assert executor_module.get_background_task_result(execution_id) is result
|
||||
assert not result.execution_done_event.is_set()
|
||||
executor_module._mark_background_execution_done(execution_id, result)
|
||||
assert executor_module.get_background_task_result(execution_id) is None
|
||||
finally:
|
||||
executor_module.force_cleanup_background_task(execution_id)
|
||||
|
||||
@ -1635,6 +1635,9 @@ def test_cleanup_called_on_timed_out(monkeypatch):
|
||||
config = _make_subagent_config()
|
||||
events = []
|
||||
cleanup_calls = []
|
||||
usage_tail_calls = []
|
||||
result = _make_result(FakeSubagentStatus.TIMED_OUT, error="timeout")
|
||||
result.execution_done_event = threading.Event()
|
||||
|
||||
monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus)
|
||||
monkeypatch.setattr(
|
||||
@ -1647,7 +1650,7 @@ def test_cleanup_called_on_timed_out(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
task_tool_module,
|
||||
"get_background_task_result",
|
||||
lambda _: _make_result(FakeSubagentStatus.TIMED_OUT, error="timeout"),
|
||||
lambda _: result,
|
||||
)
|
||||
monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: events.append)
|
||||
monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep)
|
||||
@ -1657,6 +1660,19 @@ def test_cleanup_called_on_timed_out(monkeypatch):
|
||||
"cleanup_background_task",
|
||||
lambda task_id: cleanup_calls.append(task_id),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
task_tool_module,
|
||||
"_schedule_deferred_subagent_cleanup",
|
||||
lambda runtime, execution_id, trace_id, max_polls, *, terminal_result=None: usage_tail_calls.append(
|
||||
(
|
||||
runtime,
|
||||
execution_id,
|
||||
trace_id,
|
||||
max_polls,
|
||||
terminal_result,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
output = _run_task_tool(
|
||||
runtime=_make_runtime(),
|
||||
@ -1668,6 +1684,13 @@ def test_cleanup_called_on_timed_out(monkeypatch):
|
||||
|
||||
assert _task_tool_message(output).content == "Task timed out. Error: timeout"
|
||||
assert cleanup_calls == ["tc-cleanup-timedout"]
|
||||
assert len(usage_tail_calls) == 1
|
||||
assert usage_tail_calls[0][1:] == (
|
||||
"tc-cleanup-timedout",
|
||||
"trace-1",
|
||||
task_tool_module._TERMINAL_USAGE_TAIL_POLLS,
|
||||
result,
|
||||
)
|
||||
|
||||
|
||||
def test_cleanup_not_called_on_polling_safety_timeout(monkeypatch):
|
||||
@ -2398,6 +2421,103 @@ def test_deferred_final_usage_reported_on_parent_loop_with_real_recorder(monkeyp
|
||||
parent_loop.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_terminal_usage_tail_waits_for_cleanup_fence_and_reports_on_parent_loop(monkeypatch):
|
||||
parent_loop = asyncio.get_running_loop()
|
||||
execution_done = threading.Event()
|
||||
calls: list[tuple[object, list[dict]]] = []
|
||||
scheduled: list[asyncio.Task] = []
|
||||
result = SimpleNamespace(
|
||||
status=FakeSubagentStatus.TIMED_OUT,
|
||||
usage_reported=True,
|
||||
token_usage_records=[
|
||||
{"source_run_id": "run-1", "total_tokens": 10},
|
||||
],
|
||||
execution_done_event=execution_done,
|
||||
)
|
||||
|
||||
class LoopPinnedJournal:
|
||||
def record_external_llm_usage_records(self, records):
|
||||
calls.append((asyncio.get_running_loop(), list(records)))
|
||||
|
||||
runtime = SimpleNamespace(config={"callbacks": [LoopPinnedJournal()]})
|
||||
|
||||
def schedule_locally(coro):
|
||||
task = asyncio.create_task(coro)
|
||||
scheduled.append(task)
|
||||
return task
|
||||
|
||||
cleanup_calls = []
|
||||
monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus)
|
||||
monkeypatch.setattr(
|
||||
task_tool_module,
|
||||
"run_on_isolated_subagent_loop",
|
||||
schedule_locally,
|
||||
)
|
||||
monkeypatch.setattr(task_tool_module, "cleanup_background_task", cleanup_calls.append)
|
||||
monkeypatch.setattr(task_tool_module, "_SUBAGENT_POLL_INTERVAL_SECONDS", 0.001)
|
||||
|
||||
handle = task_tool_module._schedule_deferred_subagent_cleanup(
|
||||
runtime,
|
||||
"execution-timeout",
|
||||
"trace-timeout",
|
||||
1_000,
|
||||
terminal_result=result,
|
||||
)
|
||||
assert handle is scheduled[0]
|
||||
await asyncio.sleep(0.01)
|
||||
assert calls == []
|
||||
|
||||
result.token_usage_records.append(
|
||||
{"source_run_id": "run-2", "total_tokens": 25},
|
||||
)
|
||||
execution_done.set()
|
||||
await handle
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert len(calls) == 1
|
||||
observed_loop, observed_records = calls[0]
|
||||
assert observed_loop is parent_loop
|
||||
assert [record["source_run_id"] for record in observed_records] == [
|
||||
"run-1",
|
||||
"run-2",
|
||||
]
|
||||
assert cleanup_calls == ["execution-timeout"]
|
||||
|
||||
|
||||
def test_terminal_usage_tail_reports_final_snapshot_when_cleanup_already_done(monkeypatch):
|
||||
result = _make_result(FakeSubagentStatus.TIMED_OUT, error="timeout")
|
||||
result.execution_done_event = threading.Event()
|
||||
result.execution_done_event.set()
|
||||
runtime = _make_runtime()
|
||||
report_calls = []
|
||||
scheduled = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
task_tool_module,
|
||||
"_report_subagent_usage",
|
||||
lambda runtime_arg, result_arg, **kwargs: report_calls.append(
|
||||
(runtime_arg, result_arg, kwargs),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
task_tool_module,
|
||||
"_schedule_deferred_subagent_cleanup",
|
||||
lambda *args, **kwargs: scheduled.append((args, kwargs)),
|
||||
)
|
||||
|
||||
task_tool_module._supervise_terminal_usage_tail(
|
||||
runtime,
|
||||
"execution-timeout",
|
||||
"trace-timeout",
|
||||
12,
|
||||
result,
|
||||
)
|
||||
|
||||
assert report_calls == [(runtime, result, {"final": True})]
|
||||
assert scheduled == []
|
||||
|
||||
|
||||
@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
|
||||
@ -2683,6 +2803,66 @@ def test_cancellation_wait_uses_subagent_polling_budget(monkeypatch):
|
||||
assert cleanup_calls == ["tc-cancel-budget"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_interrupted_terminal_result_defers_final_usage_until_cleanup_fence(monkeypatch):
|
||||
result = _make_result(FakeSubagentStatus.TIMED_OUT, error="timeout")
|
||||
result.execution_done_event = threading.Event()
|
||||
runtime = _make_runtime()
|
||||
report_calls = []
|
||||
cleanup_calls = []
|
||||
scheduled = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
task_tool_module,
|
||||
"get_background_task_result",
|
||||
lambda _: result,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
task_tool_module,
|
||||
"_report_subagent_usage",
|
||||
lambda runtime_arg, result_arg, **kwargs: report_calls.append(
|
||||
(runtime_arg, result_arg, kwargs),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
task_tool_module,
|
||||
"cleanup_background_task",
|
||||
cleanup_calls.append,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
task_tool_module,
|
||||
"_schedule_deferred_subagent_cleanup",
|
||||
lambda runtime_arg, execution_id, trace_id, max_polls, *, terminal_result=None: scheduled.append(
|
||||
(
|
||||
runtime_arg,
|
||||
execution_id,
|
||||
trace_id,
|
||||
max_polls,
|
||||
terminal_result,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
await task_tool_module._finalize_interrupted_subagent(
|
||||
runtime,
|
||||
"execution-timeout",
|
||||
"trace-timeout",
|
||||
12,
|
||||
)
|
||||
|
||||
assert report_calls == [(runtime, result, {})]
|
||||
assert scheduled == [
|
||||
(
|
||||
runtime,
|
||||
"execution-timeout",
|
||||
"trace-timeout",
|
||||
12,
|
||||
result,
|
||||
)
|
||||
]
|
||||
assert cleanup_calls == ["execution-timeout"]
|
||||
|
||||
|
||||
def test_cancellation_calls_request_cancel(monkeypatch):
|
||||
"""Verify CancelledError path cancels the server-generated execution ID."""
|
||||
config = _make_subagent_config()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user