mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
cd28a1cfe9
* 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>
1054 lines
39 KiB
Python
1054 lines
39 KiB
Python
import asyncio
|
|
import logging
|
|
from enum import Enum
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
|
|
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
|
from deerflow.mcp_scope import THREAD_INCARNATION_CONTEXT_KEY
|
|
from deerflow.subagents import batch_service as service_module
|
|
from deerflow.subagents.batch_runtime import BatchSubmitRequest
|
|
from deerflow.subagents.batch_service import SubagentBatchService
|
|
from deerflow.subagents.capacity import SubagentExecutionCapacity
|
|
|
|
_MISSING = object()
|
|
|
|
|
|
class FakeStatus(Enum):
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
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, FakeStatus.TIMED_OUT}
|
|
|
|
|
|
def _request(**overrides) -> BatchSubmitRequest:
|
|
values = {
|
|
"user_id": "user-1",
|
|
"thread_id": "thread-1",
|
|
"run_id": "run-1",
|
|
"tool_call_id": "call-1",
|
|
"submission_key": "run-1:call-1",
|
|
"title": "Records",
|
|
"subagent_type": "general-purpose",
|
|
"items": [{"key": "record-1", "prompt": "Process record 1"}],
|
|
"max_live_items": None,
|
|
"max_running_items": None,
|
|
"execution_spec": {
|
|
"subagent_config": {
|
|
"name": "general-purpose",
|
|
"description": "General purpose",
|
|
"system_prompt": "Work carefully.",
|
|
},
|
|
"parent_model": "model-a",
|
|
"knowledge_scope": {
|
|
"version": 1,
|
|
"mode": "selected",
|
|
"dataset_ids": ["dataset-1"],
|
|
},
|
|
},
|
|
}
|
|
values.update(overrides)
|
|
return BatchSubmitRequest(**values)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_submit_keeps_batch_running_limit_separate_from_one_process_capacity() -> None:
|
|
repository = SimpleNamespace(create_batch=AsyncMock(return_value={"id": "batch-1"}))
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(max_running_items_per_batch=32),
|
|
runtime_config=SubagentRuntimeConfig(max_running=3),
|
|
)
|
|
|
|
result = await service.submit(_request(max_live_items=20, max_running_items=10))
|
|
|
|
assert result == {"id": "batch-1"}
|
|
assert repository.create_batch.await_args.kwargs["max_running_items"] == 10
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("overrides", "expected"),
|
|
[
|
|
({"max_running_items": 0}, "max_running_items must be between 1 and 64"),
|
|
({"max_live_items": 1, "max_running_items": 0}, "max_running_items must be between 1 and 64"),
|
|
({"max_live_items": 0}, "max_live_items must be between 1 and 1000"),
|
|
],
|
|
)
|
|
async def test_submit_reports_an_explicit_zero_limit_by_name(overrides: dict[str, int], expected: str) -> None:
|
|
repository = SimpleNamespace(create_batch=AsyncMock(return_value={"id": "batch-1"}))
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(),
|
|
runtime_config=SubagentRuntimeConfig(),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match=expected):
|
|
await service.submit(_request(**overrides))
|
|
|
|
repository.create_batch.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("overrides", "live", "running"),
|
|
[
|
|
({}, 100, 3),
|
|
({"max_live_items": 40}, 40, 3),
|
|
({"max_live_items": 40, "max_running_items": 5}, 40, 5),
|
|
],
|
|
)
|
|
async def test_submit_defaults_only_the_limits_the_caller_omitted(overrides: dict[str, int], live: int, running: int) -> None:
|
|
repository = SimpleNamespace(create_batch=AsyncMock(return_value={"id": "batch-1"}))
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(),
|
|
runtime_config=SubagentRuntimeConfig(),
|
|
)
|
|
|
|
await service.submit(_request(**overrides))
|
|
|
|
kwargs = repository.create_batch.await_args.kwargs
|
|
assert (kwargs["max_live_items"], kwargs["max_running_items"]) == (live, running)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("thread_incarnation", "expected_present"),
|
|
[
|
|
("incarnation-1", True),
|
|
(None, True),
|
|
(_MISSING, False),
|
|
],
|
|
)
|
|
async def test_execute_item_marks_real_running_then_persists_terminal_result(
|
|
monkeypatch,
|
|
thread_incarnation,
|
|
expected_present,
|
|
) -> None:
|
|
result = SimpleNamespace(
|
|
status=FakeStatus.RUNNING,
|
|
result=None,
|
|
error=None,
|
|
stop_reason=None,
|
|
token_usage_records=None,
|
|
)
|
|
|
|
execution_spec = dict(_request().execution_spec)
|
|
if thread_incarnation is not _MISSING:
|
|
execution_spec[THREAD_INCARNATION_CONTEXT_KEY] = thread_incarnation
|
|
|
|
class Repository:
|
|
def __init__(self) -> None:
|
|
self.marked_running = False
|
|
self.finalized = None
|
|
|
|
async def claim_items(self, **_kwargs):
|
|
return [
|
|
{
|
|
"id": "item-1",
|
|
"item_key": "record-1",
|
|
"prompt": "Process record 1",
|
|
"batch": {
|
|
"id": "batch-1",
|
|
"thread_id": "thread-1",
|
|
"user_id": "user-1",
|
|
"run_id": "run-1",
|
|
"execution_spec": execution_spec,
|
|
},
|
|
}
|
|
]
|
|
|
|
async def mark_item_running(self, *_args, **_kwargs):
|
|
self.marked_running = True
|
|
result.status = FakeStatus.COMPLETED
|
|
result.result = "done"
|
|
return True
|
|
|
|
async def renew_item_lease(self, *_args, **_kwargs):
|
|
return {"valid": True, "cancel_requested": False}
|
|
|
|
async def finalize_item(self, *_args, **kwargs):
|
|
self.finalized = kwargs
|
|
return True
|
|
|
|
execution_capacity = SubagentExecutionCapacity(SubagentRuntimeConfig(max_running=1))
|
|
executor_kwargs = {}
|
|
|
|
class Executor:
|
|
def __init__(self, **kwargs) -> None:
|
|
executor_kwargs.update(kwargs)
|
|
|
|
def execute_async(self, _prompt, task_id=None):
|
|
assert task_id == "item-1"
|
|
return "execution-1"
|
|
|
|
repository = Repository()
|
|
monkeypatch.setattr(service_module, "get_app_config", lambda: SimpleNamespace())
|
|
monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_args, **_kwargs: "model-a")
|
|
monkeypatch.setattr(service_module, "SubagentExecutor", Executor)
|
|
monkeypatch.setattr(service_module, "SubagentStatus", FakeStatus)
|
|
monkeypatch.setattr(service_module, "get_background_task_result", lambda _execution_id: result)
|
|
monkeypatch.setattr(service_module, "cleanup_background_task", lambda _execution_id: None)
|
|
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **_kwargs: [])
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(),
|
|
runtime_config=SubagentRuntimeConfig(max_running=1),
|
|
execution_capacity=execution_capacity,
|
|
)
|
|
|
|
await service.run_once(now=service_module.datetime.now(service_module.UTC))
|
|
await asyncio.gather(*list(service._executions.values()))
|
|
|
|
assert repository.marked_running is True
|
|
assert repository.finalized is not None
|
|
assert repository.finalized["succeeded"] is True
|
|
assert repository.finalized["result"] == "done"
|
|
assert executor_kwargs["execution_capacity"] is execution_capacity
|
|
assert executor_kwargs["knowledge_scope"] == {
|
|
"version": 1,
|
|
"mode": "selected",
|
|
"dataset_ids": ["dataset-1"],
|
|
}
|
|
assert (THREAD_INCARNATION_CONTEXT_KEY in executor_kwargs) is expected_present
|
|
if expected_present:
|
|
assert executor_kwargs[THREAD_INCARNATION_CONTEXT_KEY] is thread_incarnation
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_item_polls_completion_without_waiting_for_lease_renewal(monkeypatch) -> None:
|
|
result = SimpleNamespace(
|
|
status=FakeStatus.PENDING,
|
|
result=None,
|
|
error=None,
|
|
stop_reason=None,
|
|
token_usage_records=None,
|
|
)
|
|
reads = 0
|
|
|
|
class Repository:
|
|
def __init__(self) -> None:
|
|
self.finalized = None
|
|
|
|
async def claim_items(self, **_kwargs):
|
|
return [
|
|
{
|
|
"id": "item-1",
|
|
"item_key": "record-1",
|
|
"prompt": "Process record 1",
|
|
"batch": {
|
|
"id": "batch-1",
|
|
"thread_id": "thread-1",
|
|
"user_id": "user-1",
|
|
"run_id": "run-1",
|
|
"execution_spec": _request().execution_spec,
|
|
},
|
|
}
|
|
]
|
|
|
|
async def mark_item_running(self, *_args, **_kwargs):
|
|
raise AssertionError("a task that completes between polls need not expose running")
|
|
|
|
async def renew_item_lease(self, *_args, **_kwargs):
|
|
# Renew before assembly and revalidate before dispatch; neither
|
|
# heartbeat nor polling should delay a task completing between polls.
|
|
self.renews = getattr(self, "renews", 0) + 1
|
|
return {"valid": True, "cancel_requested": False}
|
|
|
|
async def finalize_item(self, *_args, **kwargs):
|
|
self.finalized = kwargs
|
|
return True
|
|
|
|
class Executor:
|
|
def __init__(self, **_kwargs) -> None:
|
|
pass
|
|
|
|
def execute_async(self, _prompt, task_id=None):
|
|
assert task_id == "item-1"
|
|
return "execution-1"
|
|
|
|
def read_result(_execution_id):
|
|
nonlocal reads
|
|
reads += 1
|
|
if reads > 1:
|
|
result.status = FakeStatus.COMPLETED
|
|
result.result = "fast result"
|
|
return result
|
|
|
|
repository = Repository()
|
|
monkeypatch.setattr(service_module, "get_app_config", lambda: SimpleNamespace())
|
|
monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_args, **_kwargs: "model-a")
|
|
monkeypatch.setattr(service_module, "SubagentExecutor", Executor)
|
|
monkeypatch.setattr(service_module, "SubagentStatus", FakeStatus)
|
|
monkeypatch.setattr(service_module, "get_background_task_result", read_result)
|
|
monkeypatch.setattr(service_module, "cleanup_background_task", lambda _execution_id: None)
|
|
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **_kwargs: [])
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(poll_interval_seconds=0.1, lease_seconds=120),
|
|
runtime_config=SubagentRuntimeConfig(max_running=1),
|
|
)
|
|
|
|
await service.run_once(now=service_module.datetime.now(service_module.UTC))
|
|
await asyncio.wait_for(
|
|
asyncio.gather(*list(service._executions.values())),
|
|
timeout=1,
|
|
)
|
|
|
|
assert repository.finalized is not None
|
|
assert repository.finalized["result"] == "fast result"
|
|
assert repository.renews == 2
|
|
|
|
|
|
class _Registry:
|
|
"""Executor registry double with the real cleanup and cancellation timing."""
|
|
|
|
def __init__(self) -> None:
|
|
self.results: dict[str, SimpleNamespace] = {}
|
|
self.cancelled: list[str] = []
|
|
self.forced: list[str] = []
|
|
|
|
def get(self, execution_id):
|
|
return self.results.get(execution_id)
|
|
|
|
def cleanup(self, execution_id) -> None:
|
|
# The real cleanup_background_task keeps entries that are not terminal.
|
|
result = self.results.get(execution_id)
|
|
if result is not None and result.status.is_terminal:
|
|
del self.results[execution_id]
|
|
|
|
def force(self, execution_id) -> None:
|
|
self.forced.append(execution_id)
|
|
self.results.pop(execution_id, None)
|
|
|
|
def cancel(self, execution_id) -> None:
|
|
self.cancelled.append(execution_id)
|
|
result = self.results.get(execution_id)
|
|
if result is not None:
|
|
# The executor turns a cancelled run terminal on its own loop,
|
|
# after the caller has already left its finally block.
|
|
asyncio.get_running_loop().call_later(0.05, setattr, result, "status", FakeStatus.CANCELLED)
|
|
|
|
|
|
def _install_registry(monkeypatch, registry: _Registry) -> None:
|
|
monkeypatch.setattr(service_module, "SubagentStatus", FakeStatus)
|
|
monkeypatch.setattr(service_module, "get_background_task_result", registry.get)
|
|
monkeypatch.setattr(service_module, "cleanup_background_task", registry.cleanup)
|
|
monkeypatch.setattr(service_module, "force_cleanup_background_task", registry.force)
|
|
monkeypatch.setattr(service_module, "request_cancel_background_task", registry.cancel)
|
|
monkeypatch.setattr(service_module, "run_on_isolated_subagent_loop", asyncio.ensure_future)
|
|
|
|
|
|
def _running_item_service(monkeypatch, registry: _Registry, *, mark_item_running) -> tuple[SubagentBatchService, SimpleNamespace]:
|
|
registry.results["execution-1"] = SimpleNamespace(status=FakeStatus.RUNNING, result=None, error=None, stop_reason=None, token_usage_records=None)
|
|
finalized = SimpleNamespace(kwargs=None)
|
|
|
|
class Repository:
|
|
async def claim_items(self, **_kwargs):
|
|
return [
|
|
{
|
|
"id": "item-1",
|
|
"item_key": "record-1",
|
|
"prompt": "Process record 1",
|
|
"batch": {"id": "batch-1", "thread_id": "thread-1", "user_id": "user-1", "run_id": "run-1", "execution_spec": _request().execution_spec},
|
|
}
|
|
]
|
|
|
|
async def renew_item_lease(self, *_args, **_kwargs):
|
|
return {"valid": True, "cancel_requested": False}
|
|
|
|
async def mark_item_running(self, *_args, **_kwargs):
|
|
return await mark_item_running()
|
|
|
|
async def finalize_item(self, *_args, **kwargs):
|
|
finalized.kwargs = kwargs
|
|
return True
|
|
|
|
class Executor:
|
|
def __init__(self, **_kwargs) -> None:
|
|
pass
|
|
|
|
def execute_async(self, _prompt, task_id=None):
|
|
return "execution-1"
|
|
|
|
_install_registry(monkeypatch, registry)
|
|
monkeypatch.setattr(service_module, "get_app_config", lambda: SimpleNamespace())
|
|
monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_args, **_kwargs: "model-a")
|
|
monkeypatch.setattr(service_module, "SubagentExecutor", Executor)
|
|
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **_kwargs: [])
|
|
service = SubagentBatchService(
|
|
repository=Repository(),
|
|
config=SubagentBatchesConfig(poll_interval_seconds=0.1, lease_seconds=120),
|
|
runtime_config=SubagentRuntimeConfig(max_running=1),
|
|
)
|
|
return service, finalized
|
|
|
|
|
|
async def _wait_until(predicate, *, timeout: float = 2.0) -> None:
|
|
async def poll() -> None:
|
|
while not predicate():
|
|
await asyncio.sleep(0.01)
|
|
|
|
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
|
|
# finally-block cleanup alone saw a live entry, kept it, and nothing ever
|
|
# removed the SubagentResult (with its whole transcript) afterwards.
|
|
registry = _Registry()
|
|
running = asyncio.Event()
|
|
|
|
async def mark_item_running():
|
|
running.set()
|
|
return True
|
|
|
|
service, finalized = _running_item_service(monkeypatch, registry, mark_item_running=mark_item_running)
|
|
await service.run_once(now=service_module.datetime.now(service_module.UTC))
|
|
await asyncio.wait_for(running.wait(), timeout=1)
|
|
|
|
await asyncio.wait_for(service.stop(), timeout=1)
|
|
|
|
assert "execution-1" in registry.cancelled
|
|
await _wait_until(lambda: "execution-1" not in registry.results)
|
|
assert registry.forced == []
|
|
assert finalized.kwargs is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_item_error_cancels_and_releases_an_unfinished_execution(monkeypatch) -> None:
|
|
# A repository error while polling fails the durable item; the subagent
|
|
# must not keep running unsupervised, nor stay in the registry.
|
|
registry = _Registry()
|
|
|
|
async def mark_item_running():
|
|
raise RuntimeError("database is locked")
|
|
|
|
service, finalized = _running_item_service(monkeypatch, registry, mark_item_running=mark_item_running)
|
|
await service.run_once(now=service_module.datetime.now(service_module.UTC))
|
|
await asyncio.wait_for(asyncio.gather(*list(service._executions.values())), timeout=1)
|
|
|
|
assert finalized.kwargs is not None and finalized.kwargs["succeeded"] is False
|
|
assert registry.cancelled == ["execution-1"]
|
|
await _wait_until(lambda: "execution-1" not in registry.results)
|
|
assert registry.forced == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_release_removes_an_execution_that_never_turns_terminal(monkeypatch) -> None:
|
|
# Cancellation is cooperative; a run that ignores it must not keep its
|
|
# entry, or this poll, alive forever.
|
|
registry = _Registry()
|
|
registry.results["execution-1"] = SimpleNamespace(status=FakeStatus.RUNNING)
|
|
_install_registry(monkeypatch, registry)
|
|
|
|
await asyncio.wait_for(service_module._release_when_terminal("execution-1", poll_seconds=0.01, max_wait_seconds=0.05), timeout=1)
|
|
|
|
assert registry.forced == ["execution-1"]
|
|
assert registry.results == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_release_survives_an_unavailable_subagent_loop(monkeypatch) -> None:
|
|
# Runs in a finally block: failing to schedule must not replace the
|
|
# in-flight CancelledError, and must not strand an unawaited coroutine.
|
|
registry = _Registry()
|
|
registry.results["execution-1"] = SimpleNamespace(status=FakeStatus.RUNNING)
|
|
_install_registry(monkeypatch, registry)
|
|
submitted = []
|
|
|
|
def unavailable(coro):
|
|
submitted.append(coro)
|
|
raise RuntimeError("Isolated subagent event loop shutdown is still pending")
|
|
|
|
monkeypatch.setattr(service_module, "run_on_isolated_subagent_loop", unavailable)
|
|
service = SubagentBatchService(repository=SimpleNamespace(), config=SubagentBatchesConfig(), runtime_config=SubagentRuntimeConfig())
|
|
|
|
service._release_execution("execution-1")
|
|
|
|
assert registry.cancelled == ["execution-1"]
|
|
assert len(submitted) == 1 and submitted[0].cr_frame is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_executor_admission_failure_requeues_instead_of_finalizing(monkeypatch) -> None:
|
|
result = SimpleNamespace(
|
|
status=FakeStatus.FAILED,
|
|
result=None,
|
|
error="Process-wide subagent capacity is full",
|
|
stop_reason=None,
|
|
token_usage_records=None,
|
|
admission_failure=True,
|
|
)
|
|
|
|
class Repository:
|
|
def __init__(self) -> None:
|
|
self.requeued = None
|
|
self.finalized = False
|
|
|
|
async def claim_items(self, **_kwargs):
|
|
return [
|
|
{
|
|
"id": "item-1",
|
|
"item_key": "record-1",
|
|
"prompt": "Process record 1",
|
|
"batch": {
|
|
"id": "batch-1",
|
|
"thread_id": "thread-1",
|
|
"user_id": "user-1",
|
|
"run_id": "run-1",
|
|
"execution_spec": _request().execution_spec,
|
|
},
|
|
}
|
|
]
|
|
|
|
async def renew_item_lease(self, *_args, **_kwargs):
|
|
return {"valid": True, "cancel_requested": False}
|
|
|
|
async def requeue_item_after_admission_failure(self, item_id, **kwargs):
|
|
self.requeued = (item_id, kwargs)
|
|
return True
|
|
|
|
async def finalize_item(self, *_args, **_kwargs):
|
|
self.finalized = True
|
|
return True
|
|
|
|
class Executor:
|
|
def __init__(self, **_kwargs) -> None:
|
|
pass
|
|
|
|
def execute_async(self, _prompt, task_id=None):
|
|
assert task_id == "item-1"
|
|
return "execution-1"
|
|
|
|
repository = Repository()
|
|
monkeypatch.setattr(service_module, "get_app_config", lambda: SimpleNamespace())
|
|
monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_args, **_kwargs: "model-a")
|
|
monkeypatch.setattr(service_module, "SubagentExecutor", Executor)
|
|
monkeypatch.setattr(service_module, "SubagentStatus", FakeStatus)
|
|
monkeypatch.setattr(service_module, "get_background_task_result", lambda _execution_id: result)
|
|
monkeypatch.setattr(service_module, "cleanup_background_task", lambda _execution_id: None)
|
|
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **_kwargs: [])
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(),
|
|
runtime_config=SubagentRuntimeConfig(max_running=1),
|
|
)
|
|
|
|
await service.run_once(now=service_module.datetime.now(service_module.UTC))
|
|
await asyncio.gather(*list(service._executions.values()))
|
|
|
|
assert repository.requeued is not None
|
|
assert repository.requeued[0] == "item-1"
|
|
assert repository.finalized is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_during_tool_assembly_skips_launch(monkeypatch, tmp_path) -> None:
|
|
"""A batch cancelled while tool assembly is blocked must not launch.
|
|
|
|
Regression (review of 249dba82): ``_execute_item()`` called
|
|
``executor.execute_async()`` unconditionally after assembly, so
|
|
``cancel_batch()`` landing while assembly was blocked in the worker thread
|
|
terminalized the durable item but could not stop the not-yet-started
|
|
execution, and the orphaned launch still invoked the model.
|
|
"""
|
|
import threading
|
|
from datetime import UTC, datetime
|
|
|
|
from deerflow.config.database_config import DatabaseConfig
|
|
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
|
|
from deerflow.persistence.subagent_batches import SubagentBatchRepository
|
|
|
|
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
|
try:
|
|
repository = SubagentBatchRepository(get_session_factory())
|
|
await repository.create_batch(
|
|
batch_id="batch-1",
|
|
user_id="user-1",
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
tool_call_id="call-1",
|
|
submission_key="run-1:call-1",
|
|
title="Cancelled during assembly",
|
|
subagent_type="general-purpose",
|
|
items=[{"key": "record-1", "prompt": "Process record 1"}],
|
|
max_live_items=2,
|
|
max_running_items=1,
|
|
max_attempts=2,
|
|
execution_spec={
|
|
"subagent_config": {
|
|
"name": "general-purpose",
|
|
"description": "test",
|
|
"system_prompt": "sys",
|
|
},
|
|
},
|
|
)
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(lease_seconds=60),
|
|
runtime_config=SubagentRuntimeConfig(max_running=3),
|
|
app_config=SimpleNamespace(),
|
|
)
|
|
claimed = await repository.claim_items(
|
|
now=datetime.now(UTC),
|
|
lease_owner="assembly-claim",
|
|
lease_seconds=60,
|
|
limit=10,
|
|
)
|
|
assert len(claimed) == 1
|
|
item = claimed[0]
|
|
assert item["batch"]["id"] == "batch-1"
|
|
|
|
assembly_started = threading.Event()
|
|
assembly_release = threading.Event()
|
|
|
|
def _blocking_assembly(**_kwargs):
|
|
# Real blocking wait in the assembly-pool worker: parks assembly
|
|
# until the test has cancelled the batch.
|
|
assembly_started.set()
|
|
assembly_release.wait(timeout=15)
|
|
return []
|
|
|
|
monkeypatch.setattr("deerflow.tools.get_available_tools", _blocking_assembly)
|
|
monkeypatch.setattr(service_module, "SubagentStatus", FakeStatus)
|
|
monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_a, **_k: "test-model")
|
|
monkeypatch.setattr(service_module, "request_cancel_background_task", lambda _execution_id: None)
|
|
|
|
launched: list[str] = []
|
|
|
|
class _Executor:
|
|
def __init__(self, **_kwargs) -> None:
|
|
pass
|
|
|
|
def execute_async(self, _prompt, task_id=None):
|
|
launched.append(task_id or "generated")
|
|
return task_id or "generated"
|
|
|
|
monkeypatch.setattr(service_module, "SubagentExecutor", _Executor)
|
|
monkeypatch.setattr(
|
|
service_module,
|
|
"get_background_task_result",
|
|
lambda _execution_id: SimpleNamespace(
|
|
status=FakeStatus.COMPLETED,
|
|
result="done",
|
|
error=None,
|
|
stop_reason=None,
|
|
token_usage_records=[],
|
|
),
|
|
)
|
|
|
|
item_task = asyncio.create_task(service._execute_item(item, lease_owner="assembly-claim"))
|
|
assert await asyncio.to_thread(assembly_started.wait, 15)
|
|
|
|
cancelled = await service.cancel_batch(batch_id="batch-1", user_id="user-1")
|
|
assert cancelled is not None
|
|
|
|
assembly_release.set()
|
|
await asyncio.wait_for(item_task, timeout=10)
|
|
|
|
assert launched == [], "cancelled work must not launch after assembly"
|
|
batch = await repository.get_batch("batch-1", user_id="user-1")
|
|
assert batch is not None
|
|
assert batch["counts"]["cancelled"] == 1
|
|
finally:
|
|
await close_engine()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_cannot_replace_poller_while_stop_is_draining(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
service = SubagentBatchService(
|
|
repository=SimpleNamespace(),
|
|
config=SubagentBatchesConfig(),
|
|
runtime_config=SubagentRuntimeConfig(),
|
|
)
|
|
poller_entered = asyncio.Event()
|
|
allow_poller_exit = asyncio.Event()
|
|
|
|
async def blocking_poller() -> None:
|
|
poller_entered.set()
|
|
try:
|
|
await asyncio.Future()
|
|
except asyncio.CancelledError:
|
|
await allow_poller_exit.wait()
|
|
|
|
monkeypatch.setattr(service, "_run", blocking_poller)
|
|
await service.start()
|
|
original_poller = service._poller
|
|
assert original_poller is not None
|
|
await asyncio.wait_for(poller_entered.wait(), timeout=1)
|
|
|
|
stop_task = asyncio.create_task(service.stop())
|
|
await asyncio.sleep(0)
|
|
assert not stop_task.done()
|
|
assert service._poller is original_poller
|
|
with pytest.raises(RuntimeError, match="before stop completes"):
|
|
await service.start()
|
|
assert service._poller is original_poller
|
|
|
|
allow_poller_exit.set()
|
|
await asyncio.wait_for(stop_task, timeout=1)
|
|
assert original_poller.done()
|
|
assert service._poller is None
|
|
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_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
|
|
async def test_cancelled_batch_stop_allows_successful_retry(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
service = SubagentBatchService(
|
|
repository=SimpleNamespace(),
|
|
config=SubagentBatchesConfig(),
|
|
runtime_config=SubagentRuntimeConfig(),
|
|
)
|
|
poller_entered = asyncio.Event()
|
|
release_poller = asyncio.Event()
|
|
|
|
async def reluctant_poller() -> None:
|
|
poller_entered.set()
|
|
try:
|
|
await asyncio.Future()
|
|
except asyncio.CancelledError:
|
|
await release_poller.wait()
|
|
|
|
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
|
|
|
|
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()
|
|
|
|
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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_release_logs_a_failure_on_the_subagent_loop(monkeypatch, caplog) -> None:
|
|
# The release runs as a concurrent future, which never reports an
|
|
# exception nobody retrieved; a failing release would leak its entry
|
|
# without a trace.
|
|
registry = _Registry()
|
|
registry.results["execution-1"] = SimpleNamespace(status=FakeStatus.RUNNING)
|
|
_install_registry(monkeypatch, registry)
|
|
loop = asyncio.get_running_loop()
|
|
scheduled = []
|
|
|
|
def subagent_loop(coro):
|
|
scheduled.append(asyncio.run_coroutine_threadsafe(coro, loop))
|
|
return scheduled[-1]
|
|
|
|
def registry_unavailable(_execution_id):
|
|
raise RuntimeError("registry lock poisoned")
|
|
|
|
monkeypatch.setattr(service_module, "run_on_isolated_subagent_loop", subagent_loop)
|
|
service = SubagentBatchService(repository=SimpleNamespace(), config=SubagentBatchesConfig(), runtime_config=SubagentRuntimeConfig())
|
|
|
|
service._release_execution("execution-1")
|
|
monkeypatch.setattr(service_module, "get_background_task_result", registry_unavailable)
|
|
with caplog.at_level(logging.ERROR, logger=service_module.__name__):
|
|
await _wait_until(lambda: scheduled[0].done() and not service_module._pending_releases)
|
|
|
|
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")
|