mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(subagents): preserve batch stop cleanup across cancellation (#6552)
* fix(subagents): preserve batch stop cleanup across cancellation * fix(subagents): reset stop cleanup when restarting batches --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
FanouZeng-TT
GitHub
parent
5ecc1c2ac2
commit
02ce9ab283
@ -32,7 +32,7 @@ citation references, and opaque base64/URL/image sources.
|
||||
|
||||
`task` preserves the host incarnation, including null/missing/invalid values; never recapture it.
|
||||
|
||||
**Direct runtime shutdown**: `SubagentRuntime.stop()` holds its lifecycle lock until service shutdown finishes. Drain without a timeout so owned work cannot detach; propagate the first caller cancellation, chaining any service failure. Cover terminal outcomes and repeated cancellation in `tests/test_subagent_runtime.py`.
|
||||
**Direct runtime shutdown**: `SubagentRuntime.stop()` holds its lifecycle lock through an unbounded service drain; propagate first caller cancellation, chaining service failure. Test terminal outcomes/repeated cancellation in `tests/test_subagent_runtime.py`. Batch stops shield/join one cleanup generation. Restart resets its owner only after success; failed cleanup keeps startup fenced.
|
||||
|
||||
**Interrupted usage**: After stream cleanup, publish the collector's final
|
||||
snapshot on cancellation/timeout; callbacks can precede values chunks.
|
||||
@ -56,7 +56,7 @@ does not change retry policy.
|
||||
**Upload-state boundary**: Ordinary `task` delegation snapshots a valid parent `ThreadState.uploaded_files` list at dispatch, deep-copies it across the isolated-loop boundary, seeds it into the child's fresh state, and only then makes `list_uploaded_files` eligible for normal tool-policy filtering. An explicit empty list is valid and must be preserved because it means every upload in the thread is historical for this run. Missing or malformed state fails closed with the tool disabled. Durable `batch_task` execution intentionally keeps the tool disabled: delayed and recovered items have no valid parent-run upload boundary, and supporting that case requires a separate persisted-state contract.
|
||||
**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)**: Every built-in subagent execution registers `SubagentDateContextMiddleware` immediately before `SystemMessageCoalescingMiddleware`. Its one-time `before_agent` hook adds a hidden framework-owned `SystemMessage` containing only `<current_date>` before the first model call; it does not read `AppConfig.memory`, call the memory manager, rewrite the task `HumanMessage`, or inherit the lead agent's frozen-conversation/midnight lifecycle. The coalescer merges that reminder with the subagent's static prompt so strict providers still receive exactly one leading `SystemMessage`. The lead-only `DynamicContextMiddleware` registration and its date, optional-memory, and midnight-update behavior remain unchanged.
|
||||
**Execution**: Ordinary and durable-batch native subagents submit coroutines directly to one persistent isolated event loop. Shutdown keeps loop/thread ownership until the worker is confirmed stopped; a bounded join is not proof of termination, and new submissions stay fenced while cleanup is pending. Once the retained worker has exited, the next submission reaps and closes that loop under the shutdown lifecycle lock before creating a replacement. Gateway/embedded startup installs one process-wide async FIFO admission controller (default 3 running, bounded queue). Direct `create_deerflow_agent` callers can instead pass a caller-owned `SubagentRuntime`; reuse the same instance across graphs so its bound `task`, optional batch tools/service, middleware limits, and `SubagentExecutor` all share one controller without reading global YAML. An owned batch service must be started before graph construction and stopped at application shutdown. Waiters hold no scheduler thread, and cancellation/timeout release queue/slot ownership.
|
||||
**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.
|
||||
**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.
|
||||
|
||||
@ -45,6 +45,7 @@ class SubagentBatchService:
|
||||
# Also used by focused shutdown tests constructing an instance via __new__.
|
||||
_stopping: bool = False
|
||||
_stop_drains: int = 0
|
||||
_stop_cleanup_task: asyncio.Task[None] | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@ -69,6 +70,7 @@ class SubagentBatchService:
|
||||
self._poller: asyncio.Task[None] | None = None
|
||||
self._stopping = False
|
||||
self._stop_drains = 0
|
||||
self._stop_cleanup_task: asyncio.Task[None] | None = None
|
||||
self._executions: dict[str, asyncio.Task[None]] = {}
|
||||
self._execution_ids: dict[str, str] = {}
|
||||
self._item_batches: dict[str, str] = {}
|
||||
@ -76,51 +78,64 @@ class SubagentBatchService:
|
||||
async def start(self) -> None:
|
||||
if self._stopping:
|
||||
raise RuntimeError("cannot start subagent batch poller before stop completes")
|
||||
cleanup_task = getattr(self, "_stop_cleanup_task", None)
|
||||
if cleanup_task is not None:
|
||||
cleanup_task.result()
|
||||
self._stop_cleanup_task = None
|
||||
if self._poller is not None:
|
||||
return
|
||||
self._stop.clear()
|
||||
self._poller = asyncio.create_task(self._run(), name="subagent-batch-poller")
|
||||
|
||||
async def stop(self) -> None:
|
||||
# Keep poller ownership visible until the entire drain finishes.
|
||||
# Otherwise start() could create a fresh poller during this await.
|
||||
self._stop_drains += 1
|
||||
self._stopping = True
|
||||
self._stop.set()
|
||||
poller = self._poller
|
||||
cleanup_task = getattr(self, "_stop_cleanup_task", None)
|
||||
if cleanup_task is None:
|
||||
# Fence start() before taking the owned-work snapshot. All stop
|
||||
# callers then join this one cleanup generation instead of issuing
|
||||
# competing cancellations and clearing one another's state.
|
||||
self._stopping = True
|
||||
self._stop.set()
|
||||
poller = self._poller
|
||||
|
||||
# Issue every owned-work cancellation before the first await. The
|
||||
# Gateway wraps this stop hook in a deadline; if poller teardown is
|
||||
# slow, cancellation of this coroutine must not prevent native/item
|
||||
# cancellation from being requested.
|
||||
execution_ids = list(self._execution_ids.values())
|
||||
for execution_id in execution_ids:
|
||||
request_cancel_background_task(execution_id)
|
||||
tasks = list(self._executions.values())
|
||||
# Issue every owned-work cancellation before the first await. The
|
||||
# Gateway wraps this stop hook in a deadline; if poller teardown is
|
||||
# slow, cancellation of this caller must not prevent native/item
|
||||
# cancellation from being requested.
|
||||
execution_ids = list(self._execution_ids.values())
|
||||
for execution_id in execution_ids:
|
||||
request_cancel_background_task(execution_id)
|
||||
tasks = list(self._executions.values())
|
||||
|
||||
if poller is not None:
|
||||
poller.cancel()
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
|
||||
completed = False
|
||||
try:
|
||||
if poller is not None:
|
||||
await asyncio.gather(poller, return_exceptions=True)
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
self._executions.clear()
|
||||
self._execution_ids.clear()
|
||||
self._item_batches.clear()
|
||||
if self._poller is poller:
|
||||
self._poller = None
|
||||
completed = True
|
||||
finally:
|
||||
# A cancelled caller cannot leave a phantom drain count. Keep
|
||||
# start fenced until a subsequent successful stop retries cleanup.
|
||||
self._stop_drains -= 1
|
||||
if completed and self._stop_drains == 0:
|
||||
self._stopping = False
|
||||
poller.cancel()
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
|
||||
self._stop_drains = 1
|
||||
cleanup_task = asyncio.create_task(
|
||||
self._drain_stop(poller, tasks),
|
||||
name="subagent-batch-stop-cleanup",
|
||||
)
|
||||
self._stop_cleanup_task = cleanup_task
|
||||
|
||||
await asyncio.shield(cleanup_task)
|
||||
|
||||
async def _drain_stop(
|
||||
self,
|
||||
poller: asyncio.Task[None] | None,
|
||||
tasks: list[asyncio.Task[None]],
|
||||
) -> None:
|
||||
if poller is not None:
|
||||
await asyncio.gather(poller, return_exceptions=True)
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
self._executions.clear()
|
||||
self._execution_ids.clear()
|
||||
self._item_batches.clear()
|
||||
if self._poller is poller:
|
||||
self._poller = None
|
||||
self._stop_drains = 0
|
||||
self._stopping = False
|
||||
|
||||
async def _run(self) -> None:
|
||||
while not self._stop.is_set():
|
||||
|
||||
@ -530,8 +530,139 @@ async def test_start_cannot_replace_poller_while_stop_is_draining(monkeypatch: p
|
||||
assert not service._stopping
|
||||
|
||||
await service.start()
|
||||
assert service._poller is not original_poller
|
||||
await service.stop()
|
||||
restarted_poller = service._poller
|
||||
assert restarted_poller is not None
|
||||
assert restarted_poller is not original_poller
|
||||
await asyncio.wait_for(service.stop(), timeout=1)
|
||||
assert restarted_poller.done()
|
||||
assert service._poller is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restarted_batch_stop_cancels_and_drains_owned_work(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
service = SubagentBatchService(
|
||||
repository=SimpleNamespace(),
|
||||
config=SubagentBatchesConfig(),
|
||||
runtime_config=SubagentRuntimeConfig(),
|
||||
)
|
||||
poller_entered = asyncio.Event()
|
||||
item_entered = asyncio.Event()
|
||||
requested: list[str] = []
|
||||
owned_tasks: list[asyncio.Task[None]] = []
|
||||
|
||||
async def poller() -> None:
|
||||
poller_entered.set()
|
||||
await asyncio.Future()
|
||||
|
||||
async def item_work() -> None:
|
||||
item_entered.set()
|
||||
await asyncio.Future()
|
||||
|
||||
monkeypatch.setattr(service, "_run", poller)
|
||||
monkeypatch.setattr(service_module, "request_cancel_background_task", requested.append)
|
||||
|
||||
try:
|
||||
await service.start()
|
||||
first_poller = service._poller
|
||||
assert first_poller is not None
|
||||
owned_tasks.append(first_poller)
|
||||
await asyncio.wait_for(poller_entered.wait(), timeout=1)
|
||||
await asyncio.wait_for(service.stop(), timeout=1)
|
||||
first_cleanup = service._stop_cleanup_task
|
||||
assert first_poller.done()
|
||||
|
||||
poller_entered.clear()
|
||||
await service.start()
|
||||
restarted_poller = service._poller
|
||||
assert restarted_poller is not None
|
||||
assert restarted_poller is not first_poller
|
||||
owned_tasks.append(restarted_poller)
|
||||
await asyncio.wait_for(poller_entered.wait(), timeout=1)
|
||||
|
||||
item_task = asyncio.create_task(item_work())
|
||||
owned_tasks.append(item_task)
|
||||
service._executions["item-2"] = item_task
|
||||
service._execution_ids["item-2"] = "execution-2"
|
||||
service._item_batches["item-2"] = "batch-2"
|
||||
await asyncio.wait_for(item_entered.wait(), timeout=1)
|
||||
|
||||
await asyncio.wait_for(service.stop(), timeout=1)
|
||||
assert requested == ["execution-2"]
|
||||
assert restarted_poller.done()
|
||||
assert item_task.done()
|
||||
assert service._poller is None
|
||||
assert service._executions == {}
|
||||
assert service._execution_ids == {}
|
||||
assert service._item_batches == {}
|
||||
assert service._stop.is_set()
|
||||
assert not service._stopping
|
||||
second_cleanup = service._stop_cleanup_task
|
||||
assert second_cleanup is not None
|
||||
assert second_cleanup is not first_cleanup
|
||||
assert second_cleanup.done()
|
||||
|
||||
await asyncio.wait_for(service.stop(), timeout=1)
|
||||
assert service._stop_cleanup_task is second_cleanup
|
||||
assert requested == ["execution-2"]
|
||||
finally:
|
||||
for task in owned_tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(*owned_tasks, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_batch_stops_share_one_cleanup(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
service = SubagentBatchService(
|
||||
repository=SimpleNamespace(),
|
||||
config=SubagentBatchesConfig(),
|
||||
runtime_config=SubagentRuntimeConfig(),
|
||||
)
|
||||
poller_entered = asyncio.Event()
|
||||
poller_cancelled = asyncio.Event()
|
||||
release_poller = asyncio.Event()
|
||||
drain_calls = 0
|
||||
|
||||
async def reluctant_poller() -> None:
|
||||
poller_entered.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
except asyncio.CancelledError:
|
||||
poller_cancelled.set()
|
||||
await release_poller.wait()
|
||||
|
||||
original_drain = service._drain_stop
|
||||
|
||||
async def counted_drain(poller, tasks) -> None:
|
||||
nonlocal drain_calls
|
||||
drain_calls += 1
|
||||
await original_drain(poller, tasks)
|
||||
|
||||
monkeypatch.setattr(service, "_run", reluctant_poller)
|
||||
monkeypatch.setattr(service, "_drain_stop", counted_drain)
|
||||
await service.start()
|
||||
await asyncio.wait_for(poller_entered.wait(), timeout=1)
|
||||
|
||||
first_stop = asyncio.create_task(service.stop())
|
||||
await asyncio.wait_for(poller_cancelled.wait(), timeout=1)
|
||||
cleanup_task = service._stop_cleanup_task
|
||||
second_stop = asyncio.create_task(service.stop())
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert cleanup_task is not None
|
||||
assert service._stop_cleanup_task is cleanup_task
|
||||
assert service._stop_drains == 1
|
||||
assert drain_calls == 1
|
||||
assert not first_stop.done()
|
||||
assert not second_stop.done()
|
||||
|
||||
release_poller.set()
|
||||
await asyncio.wait_for(asyncio.gather(first_stop, second_stop), timeout=1)
|
||||
assert service._stop_cleanup_task is cleanup_task
|
||||
assert cleanup_task.done()
|
||||
assert service._stop_drains == 0
|
||||
assert not service._stopping
|
||||
assert drain_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -553,19 +684,92 @@ async def test_cancelled_batch_stop_allows_successful_retry(monkeypatch: pytest.
|
||||
|
||||
monkeypatch.setattr(service, "_run", reluctant_poller)
|
||||
await service.start()
|
||||
original_poller = service._poller
|
||||
await asyncio.wait_for(poller_entered.wait(), timeout=1)
|
||||
|
||||
first_stop = asyncio.create_task(service.stop())
|
||||
await asyncio.sleep(0)
|
||||
cleanup_task = service._stop_cleanup_task
|
||||
first_stop.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await first_stop
|
||||
|
||||
assert cleanup_task is not None
|
||||
assert service._stop_cleanup_task is cleanup_task
|
||||
assert service._stop_drains == 1
|
||||
assert service._stopping
|
||||
assert original_poller is not None
|
||||
assert not original_poller.done()
|
||||
assert not original_poller.cancelled()
|
||||
with pytest.raises(RuntimeError, match="before stop completes"):
|
||||
await service.start()
|
||||
|
||||
retry_stop = asyncio.create_task(service.stop())
|
||||
await asyncio.sleep(0)
|
||||
assert service._stop_cleanup_task is cleanup_task
|
||||
assert not retry_stop.done()
|
||||
|
||||
release_poller.set()
|
||||
await asyncio.wait_for(retry_stop, timeout=1)
|
||||
assert original_poller.done()
|
||||
assert service._poller is None
|
||||
assert service._stop_cleanup_task is cleanup_task
|
||||
assert cleanup_task.done()
|
||||
assert service._stop_drains == 0
|
||||
assert not service._stopping
|
||||
|
||||
await service.start()
|
||||
restarted_poller = service._poller
|
||||
assert restarted_poller is not None
|
||||
assert restarted_poller is not original_poller
|
||||
await asyncio.wait_for(service.stop(), timeout=1)
|
||||
assert restarted_poller.done()
|
||||
assert service._poller is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_stop_cleanup_failure_is_reported_to_retry(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
service = SubagentBatchService(
|
||||
repository=SimpleNamespace(),
|
||||
config=SubagentBatchesConfig(),
|
||||
runtime_config=SubagentRuntimeConfig(),
|
||||
)
|
||||
poller_entered = asyncio.Event()
|
||||
release_poller = asyncio.Event()
|
||||
cleanup_error = RuntimeError("stop cleanup failed")
|
||||
|
||||
async def reluctant_poller() -> None:
|
||||
poller_entered.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
except asyncio.CancelledError:
|
||||
await release_poller.wait()
|
||||
|
||||
class FailingClearDict(dict):
|
||||
def clear(self) -> None:
|
||||
raise cleanup_error
|
||||
|
||||
monkeypatch.setattr(service, "_run", reluctant_poller)
|
||||
service._executions = FailingClearDict()
|
||||
await service.start()
|
||||
await asyncio.wait_for(poller_entered.wait(), timeout=1)
|
||||
first_stop = asyncio.create_task(service.stop())
|
||||
await asyncio.sleep(0)
|
||||
first_stop.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await first_stop
|
||||
assert service._stop_drains == 0
|
||||
|
||||
cleanup_task = service._stop_cleanup_task
|
||||
assert cleanup_task is not None
|
||||
release_poller.set()
|
||||
with pytest.raises(RuntimeError, match="stop cleanup failed"):
|
||||
await asyncio.wait_for(asyncio.shield(cleanup_task), timeout=1)
|
||||
|
||||
assert service._stopping
|
||||
with pytest.raises(RuntimeError, match="before stop completes"):
|
||||
await service.start()
|
||||
|
||||
release_poller.set()
|
||||
await asyncio.wait_for(service.stop(), timeout=1)
|
||||
assert service._poller is None
|
||||
assert service._stop_drains == 0
|
||||
assert not service._stopping
|
||||
with pytest.raises(RuntimeError, match="stop cleanup failed") as raised:
|
||||
await service.stop()
|
||||
assert raised.value is cleanup_error
|
||||
assert service._stop_cleanup_task is cleanup_task
|
||||
|
||||
@ -3,12 +3,15 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import inspect
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
|
||||
from app.gateway import app as gateway_app
|
||||
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
|
||||
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
||||
from deerflow.subagents import batch_service as batch_service_module
|
||||
from deerflow.subagents.batch_service import SubagentBatchService
|
||||
|
||||
@ -39,6 +42,73 @@ async def test_subagent_batch_service_shutdown_is_bounded(
|
||||
assert "Subagent batch service shutdown exceeded" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gateway_shutdown_timeout_does_not_cancel_batch_cleanup(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
poller_entered = asyncio.Event()
|
||||
poller_cancelled = asyncio.Event()
|
||||
release_poller = asyncio.Event()
|
||||
|
||||
async def reluctant_poller() -> None:
|
||||
poller_entered.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
except asyncio.CancelledError:
|
||||
poller_cancelled.set()
|
||||
await release_poller.wait()
|
||||
|
||||
service = SubagentBatchService(
|
||||
repository=SimpleNamespace(),
|
||||
config=SubagentBatchesConfig(),
|
||||
runtime_config=SubagentRuntimeConfig(),
|
||||
)
|
||||
monkeypatch.setattr(service, "_run", reluctant_poller)
|
||||
await service.start()
|
||||
original_poller = service._poller
|
||||
await asyncio.wait_for(poller_entered.wait(), timeout=1)
|
||||
|
||||
app = FastAPI()
|
||||
app.state.subagent_batch_service = service
|
||||
app.state.subagent_batches_available = True
|
||||
monkeypatch.setattr(gateway_app, "_SHUTDOWN_HOOK_TIMEOUT_SECONDS", 0.01)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="app.gateway.app"):
|
||||
await gateway_app._shutdown_subagent_batch_service(app)
|
||||
await asyncio.wait_for(poller_cancelled.wait(), timeout=1)
|
||||
|
||||
cleanup_task = service._stop_cleanup_task
|
||||
assert cleanup_task is not None
|
||||
assert not cleanup_task.done()
|
||||
assert original_poller is not None
|
||||
assert not original_poller.done()
|
||||
assert not original_poller.cancelled()
|
||||
assert service._stopping
|
||||
with pytest.raises(RuntimeError, match="before stop completes"):
|
||||
await service.start()
|
||||
|
||||
retry_stop = asyncio.create_task(service.stop())
|
||||
await asyncio.sleep(0)
|
||||
assert service._stop_cleanup_task is cleanup_task
|
||||
release_poller.set()
|
||||
await asyncio.wait_for(retry_stop, timeout=1)
|
||||
|
||||
assert service._poller is None
|
||||
assert service._stop_cleanup_task is cleanup_task
|
||||
assert cleanup_task.done()
|
||||
assert not service._stopping
|
||||
assert "Subagent batch service shutdown exceeded" in caplog.text
|
||||
|
||||
await service.start()
|
||||
restarted_poller = service._poller
|
||||
assert restarted_poller is not None
|
||||
assert restarted_poller is not original_poller
|
||||
await asyncio.wait_for(service.stop(), timeout=1)
|
||||
assert restarted_poller.done()
|
||||
assert service._poller is None
|
||||
|
||||
|
||||
def test_lifespan_uses_bounded_subagent_batch_shutdown() -> None:
|
||||
source = inspect.getsource(gateway_app.lifespan)
|
||||
assert "await _shutdown_subagent_batch_service(app)" in source
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user