fix(subagents): harden background-task registry and capacity snapshot edge cases (#5086)

* fix(subagents): harden background-task registry and capacity snapshot edge cases

- execute_async drops the just-registered background entry when submitting
  to the isolated loop fails. The caller sees the exception and never
  polls, and cleanup_background_task refuses non-terminal entries, so the
  entry would otherwise stay as a PENDING zombie forever.
- SubagentExecutionCapacity.snapshot derives queued from len(_waiters)
  instead of iterating it. snapshot is read from non-loop threads (e.g.
  configure_subagent_execution_capacity) while the loop thread mutates the
  deque, so iteration can raise 'deque mutated during iteration'. The raw
  length may count a waiter that just timed out but has not removed itself
  yet, which only makes the busy-check more conservative.

* fix(subagents): close failure-path gaps around background submit

Address both review findings on the background submission lifecycle:

- execute_async() copies the isolated-loop context before registering
  the _background_tasks entry, so a context-copy failure (callback-
  manager copy or loop-bound handler filtering) can no longer strand a
  permanent PENDING entry the caller will never poll.
- _submit_to_isolated_loop_in_context() resolves the loop before
  calling the coroutine factory. As direct run_coroutine_threadsafe
  arguments the coroutine was created first, so a loop-startup failure
  stranded a never-awaited coroutine (RuntimeWarning + retained
  captures until collection). Both call sites share the fix.

New tests verified red on the previous implementation, green after:
- context-copy failure leaves no registry residue
- the real submit helper (only the loop getter patched) never invokes
  the coroutine factory when loop startup fails

* fix(subagents): close the coroutine when scheduling rejects it

run_coroutine_threadsafe can itself raise once the coroutine exists (e.g.
the loop closes between the lookup and the internal call_soon_threadsafe).
Wrap the call, close the rejected coroutine, and re-raise; a focused test
patches only run_coroutine_threadsafe and asserts the created coroutine
reaches CORO_CLOSED.
This commit is contained in:
hataa 2026-08-30 10:43:58 +08:00 committed by GitHub
parent 8a830f6354
commit 2f8d1cfc21
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 231 additions and 11 deletions

View File

@ -43,11 +43,18 @@ class SubagentExecutionCapacity:
self._waiters: deque[asyncio.Future[None]] = deque()
def snapshot(self) -> SubagentCapacitySnapshot:
# snapshot() is read from non-loop threads (notably
# configure_subagent_execution_capacity), while the loop thread mutates
# the waiters deque under its own asyncio lock — iterating cross-thread
# can raise "deque mutated during iteration". ``len()`` reads the
# deque's size atomically instead. The raw length may count a waiter
# that just timed out but has not removed itself yet; that only makes
# the "capacity busy" answer more conservative, never less.
return SubagentCapacitySnapshot(
max_running=self._config.max_running,
running=self._running,
max_queued=self._config.max_queued,
queued=sum(not waiter.done() for waiter in self._waiters),
queued=len(self._waiters),
admission_policy=self._config.admission_policy,
)

View File

@ -435,13 +435,34 @@ def _submit_to_isolated_loop_in_context(
context: Context,
coro_factory: Callable[[], Coroutine[Any, Any, SubagentResult]],
) -> Future[SubagentResult]:
"""Submit a coroutine to the isolated loop while preserving ContextVar state."""
return context.run(
lambda: asyncio.run_coroutine_threadsafe(
coro_factory(),
_get_isolated_subagent_loop(),
)
)
"""Submit a coroutine to the isolated loop while preserving ContextVar state.
The loop must be resolved before the coroutine is created: as direct
``run_coroutine_threadsafe(coro_factory(), ...)`` arguments, Python
evaluates the coroutine first, so a loop-startup failure would strand a
created-but-never-scheduled coroutine (``RuntimeWarning: coroutine ...
was never awaited``) holding its captures until collection.
Scheduling itself can still reject an already-created coroutine e.g.
the loop closes between the lookup above and the ``call_soon_threadsafe``
inside ``run_coroutine_threadsafe`` so a rejected coroutine is closed
before the error propagates.
"""
def _submit() -> Future[SubagentResult]:
loop = _get_isolated_subagent_loop()
coroutine = coro_factory()
try:
return asyncio.run_coroutine_threadsafe(coroutine, loop)
except BaseException:
# run_coroutine_threadsafe has no cleanup path for this window.
# The coroutine has not started (CORO_CREATED), so close() cannot
# run any of its body — it only releases the object and its
# captures instead of leaving them until collection.
coroutine.close()
raise
return context.run(_submit)
def _copy_isolated_subagent_context() -> Context:
@ -1441,11 +1462,16 @@ class SubagentExecutor:
self.config.timeout_seconds,
)
# Copy the parent context before registering: context copying can
# itself fail (callback-manager copy or loop-bound handler filtering),
# and a failure after registration would strand a PENDING entry —
# the caller never receives an execution_id to poll, and
# cleanup_background_task() refuses non-terminal entries.
parent_context = _copy_isolated_subagent_context()
with _background_tasks_lock:
_background_tasks[execution_id] = result
parent_context = _copy_isolated_subagent_context()
async def run_with_timeout() -> SubagentResult:
try:
return await asyncio.wait_for(
@ -1473,7 +1499,18 @@ class SubagentExecutor:
result.try_set_terminal(SubagentStatus.FAILED, error=str(exc))
return result
execution_future = _submit_to_isolated_loop_in_context(parent_context, run_with_timeout)
try:
execution_future = _submit_to_isolated_loop_in_context(parent_context, run_with_timeout)
except Exception:
# Submitting can fail before any coroutine starts (e.g. the
# persistent loop failed to spin up). The caller then sees the
# exception and never polls this execution_id, and
# cleanup_background_task() refuses non-terminal entries — so the
# just-registered entry must be dropped here, not left as a
# PENDING zombie nothing will ever remove.
with _background_tasks_lock:
_background_tasks.pop(execution_id, None)
raise
with _background_tasks_lock:
_background_futures[execution_id] = execution_future

View File

@ -0,0 +1,59 @@
"""Cross-thread safety of the shared subagent execution capacity."""
import asyncio
from collections import deque
import pytest
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
from deerflow.subagents.capacity import SubagentExecutionCapacity
class _NoIterationDeque(deque):
"""A deque whose iteration always fails, as if mutated concurrently."""
def __iter__(self):
raise RuntimeError("deque mutated during iteration")
def test_snapshot_does_not_iterate_waiters():
"""snapshot() must be readable from a non-loop thread.
``configure_subagent_execution_capacity`` reads a snapshot while the loop
thread owns the waiters deque; iterating it cross-thread can raise
``deque mutated during iteration``. The snapshot must therefore derive
``queued`` without iterating.
"""
capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=2, max_queued=4))
capacity._waiters = _NoIterationDeque(range(2))
snapshot = capacity.snapshot()
assert snapshot.max_running == 2
assert snapshot.max_queued == 4
assert snapshot.queued == 2
assert snapshot.running == 0
assert snapshot.admission_policy == "queue"
@pytest.mark.asyncio
async def test_snapshot_reports_running_and_queued_waiters():
"""The derived queued count still reflects pending waiters."""
capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1, max_queued=4, queue_timeout_seconds=5))
async def hold_slot():
async with capacity.slot():
await asyncio.sleep(60)
holder = asyncio.create_task(hold_slot())
queued = asyncio.create_task(capacity._acquire())
try:
await asyncio.sleep(0.05)
snapshot = capacity.snapshot()
assert snapshot.running == 1
assert snapshot.queued == 1
finally:
queued.cancel()
holder.cancel()
await asyncio.gather(queued, holder, return_exceptions=True)

View File

@ -16,6 +16,7 @@ the real implementation in isolation.
import asyncio
import importlib
import inspect
import sys
import threading
import time
@ -2082,6 +2083,122 @@ class TestCleanupBackgroundTask:
return _patch_default_get_app_config(importlib.reload(executor))
def test_execute_async_removes_entry_when_submit_fails(self, executor_module, classes, base_config):
"""A failed submit must not leave a PENDING entry nothing will ever poll.
The registry entry is created before the coroutine is submitted to the
isolated loop. When submission itself raises (e.g. the loop failed to
start), the caller sees the exception and never polls and
``cleanup_background_task`` refuses non-terminal entries so the fix
must drop the entry on the submit-failure path.
"""
SubagentExecutor = classes["SubagentExecutor"]
executor = SubagentExecutor(
config=base_config,
tools=[],
thread_id="test-thread",
trace_id="submit-failure-trace",
)
def failing_submit(_context, _coro_factory):
raise RuntimeError("isolated subagent event loop failed to start")
with patch.object(executor_module, "_submit_to_isolated_loop_in_context", side_effect=failing_submit):
with pytest.raises(RuntimeError, match="isolated subagent event loop"):
executor.execute_async("Task")
leftovers = [r for r in executor_module.list_background_tasks() if r.trace_id == "submit-failure-trace"]
assert leftovers == []
def test_execute_async_registers_nothing_when_context_copy_fails(self, executor_module, classes, base_config):
"""A context-copy failure must not leave a PENDING entry either.
``_copy_isolated_subagent_context`` (callback-manager copy or
loop-bound handler filtering) can raise before the coroutine is ever
submitted. The registry entry must not exist at that point yet
the caller gets no execution_id to poll, and
``cleanup_background_task`` refuses non-terminal entries, so a
registration before the copy would strand the same permanent
PENDING entry the submit-failure path already guards against.
"""
SubagentExecutor = classes["SubagentExecutor"]
executor = SubagentExecutor(
config=base_config,
tools=[],
thread_id="test-thread",
trace_id="context-copy-failure-trace",
)
def failing_context_copy():
raise RuntimeError("callback manager copy failed")
with patch.object(executor_module, "_copy_isolated_subagent_context", side_effect=failing_context_copy):
with pytest.raises(RuntimeError, match="callback manager copy"):
executor.execute_async("Task")
leftovers = [r for r in executor_module.list_background_tasks() if r.trace_id == "context-copy-failure-trace"]
assert leftovers == []
def test_submit_helper_skips_coroutine_creation_when_loop_startup_fails(self, executor_module):
"""Loop-startup failure must not strand an unscheduled coroutine.
Exercises the real ``_submit_to_isolated_loop_in_context`` (only the
loop getter is patched) rather than mocking the whole helper: the
loop must be resolved before the coroutine is created. If the
coroutine factory ran first, the created coroutine would be neither
scheduled nor closed ``RuntimeWarning: coroutine ... was never
awaited`` retaining its captures until collection.
"""
factory_calls = []
def coro_factory():
factory_calls.append("created")
async def never_scheduled(): # pragma: no cover - must not run
return None
return never_scheduled()
def failing_loop():
raise RuntimeError("Timed out starting isolated subagent event loop")
with patch.object(executor_module, "_get_isolated_subagent_loop", side_effect=failing_loop):
with pytest.raises(RuntimeError, match="Timed out starting"):
executor_module._submit_to_isolated_loop_in_context(executor_module.copy_context(), coro_factory)
assert factory_calls == []
def test_submit_helper_closes_coroutine_when_scheduling_rejects_it(self, executor_module):
"""Scheduling rejection after creation must close the coroutine.
Resolving the loop before calling the factory covers loop-startup
failure, but ``run_coroutine_threadsafe`` can itself raise once the
coroutine exists (e.g. the loop closes between the lookup and the
internal ``call_soon_threadsafe``). Only ``run_coroutine_threadsafe``
is patched: the helper must close the rejected coroutine otherwise
it stays in ``CORO_CREATED`` and re-triggers the never-awaited
warning and retained captures the startup fix already guards against.
"""
created = []
def coro_factory():
async def pending():
return None # pragma: no cover - must never run
coroutine = pending()
created.append(coroutine)
return coroutine
with patch.object(executor_module, "_get_isolated_subagent_loop", return_value=object()):
with patch.object(executor_module.asyncio, "run_coroutine_threadsafe", side_effect=RuntimeError("Event loop is closed")):
with pytest.raises(RuntimeError, match="Event loop is closed"):
executor_module._submit_to_isolated_loop_in_context(executor_module.copy_context(), coro_factory)
assert len(created) == 1
assert inspect.getcoroutinestate(created[0]) is inspect.CORO_CLOSED
def test_cleanup_removes_terminal_completed_task(self, executor_module, classes):
"""Test that cleanup removes a COMPLETED task."""
SubagentResult = classes["SubagentResult"]