mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
ce0fa4c114
* feat(batch): add browser review for persisted results * fix(batch): address browser review feedback * test(batch): cover persisted result review regressions Exercise running-to-terminal detail refresh, scoped citation dialogs, omitted evidence, translated page labels, and legacy evidence bounds. Fix component lint and formatting and synchronize the review contract documentation. * docs(frontend): keep batch review guidance within budget * fix(batch): refresh retried reviews and account for omitted evidence * fix(ci): keep batch review guidance within size budget --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
1217 lines
46 KiB
Python
1217 lines
46 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_read_browser_batch_item_projects_report_and_captured_evidence() -> None:
|
|
repository = SimpleNamespace(
|
|
get_batch=AsyncMock(return_value={"id": "batch-1", "thread_id": "thread-1"}),
|
|
list_items=AsyncMock(
|
|
return_value=[
|
|
{
|
|
"id": "item-1",
|
|
"item_key": "record-1",
|
|
"position": 0,
|
|
"status": "succeeded",
|
|
"attempt": 1,
|
|
"model_name": "model-a",
|
|
"result": "A report",
|
|
"result_truncated": False,
|
|
"error": None,
|
|
"stop_reason": None,
|
|
"acceptance_criteria": ["file:answer.md exists"],
|
|
"acceptance_verdict": {"status": "passed"},
|
|
"result_artifact": {
|
|
"knowledge_sources": {
|
|
"version": 1,
|
|
"sources": [
|
|
{
|
|
"id": "a" * 32 + "-1",
|
|
"provider": "ragflow",
|
|
"dataset_id": "dataset-1",
|
|
"document_id": "document-1",
|
|
"chunk_id": "chunk-1",
|
|
"dataset_name": "Private KB",
|
|
"document_name": "report.md",
|
|
"text": "Captured excerpt",
|
|
"pages": [2, True, 0, 3],
|
|
"truncated": True,
|
|
"secret": "must not escape",
|
|
},
|
|
{"provider": "other", "text": "drop me"},
|
|
],
|
|
"omitted_count": 1,
|
|
}
|
|
},
|
|
"started_at": "2026-10-17T00:00:00Z",
|
|
"completed_at": "2026-10-17T00:01:00Z",
|
|
"updated_at": "2026-10-17T00:01:00Z",
|
|
}
|
|
]
|
|
),
|
|
)
|
|
service = SubagentBatchService(
|
|
repository=repository,
|
|
config=SubagentBatchesConfig(),
|
|
runtime_config=SubagentRuntimeConfig(),
|
|
)
|
|
|
|
detail = await service.read_browser_batch_item(
|
|
batch_id="batch-1",
|
|
user_id="user-1",
|
|
thread_id="thread-1",
|
|
position=0,
|
|
)
|
|
|
|
assert detail["report"] == "A report"
|
|
assert detail["acceptance_verdict"] == {"status": "passed"}
|
|
assert detail["evidence"] == {
|
|
"version": 1,
|
|
"sources": [
|
|
{
|
|
"id": "a" * 32 + "-1",
|
|
"provider": "ragflow",
|
|
"dataset_id": "dataset-1",
|
|
"document_id": "document-1",
|
|
"chunk_id": "chunk-1",
|
|
"dataset_name": "Private KB",
|
|
"document_name": "report.md",
|
|
"text": "Captured excerpt",
|
|
"pages": [2, 3],
|
|
"truncated": True,
|
|
}
|
|
],
|
|
"omitted_count": 1,
|
|
}
|
|
assert "secret" not in detail["evidence"]["sources"][0]
|
|
repository.list_items.assert_awaited_once_with(
|
|
"batch-1",
|
|
user_id="user-1",
|
|
offset=0,
|
|
limit=1,
|
|
include_result=True,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("field,limit", [("id", 256), ("dataset_id", 256), ("document_id", 256), ("chunk_id", 256), ("dataset_name", 512), ("document_name", 512), ("text", 16_000)])
|
|
@pytest.mark.parametrize("oversized", [False, True])
|
|
async def test_browser_evidence_bounds_legacy_rows(field, limit, oversized):
|
|
source = {
|
|
"id": "a" * 32 + "-1",
|
|
"provider": "ragflow",
|
|
"dataset_id": "dataset-1",
|
|
"document_id": "document-1",
|
|
"chunk_id": "chunk-1",
|
|
"dataset_name": "Private KB",
|
|
"document_name": "report.md",
|
|
"text": "Captured excerpt",
|
|
"pages": [2],
|
|
"truncated": False,
|
|
}
|
|
source[field] = "x" * (limit + int(oversized))
|
|
repository = SimpleNamespace(
|
|
get_batch=AsyncMock(return_value={"id": "batch-1", "thread_id": "thread-1"}),
|
|
list_items=AsyncMock(
|
|
return_value=[
|
|
{
|
|
"position": 0,
|
|
"result": "Stored report",
|
|
"prompt": "Private prompt",
|
|
"execution_spec": {"api_key": "private-key"},
|
|
"result_artifact": {"knowledge_sources": {"version": 1, "sources": [source], "omitted_count": 2}},
|
|
}
|
|
]
|
|
),
|
|
)
|
|
service = SubagentBatchService(repository=repository, config=SubagentBatchesConfig(), runtime_config=SubagentRuntimeConfig())
|
|
detail = await service.read_browser_batch_item(batch_id="batch-1", user_id="user-1", thread_id="thread-1", position=0)
|
|
assert detail["report"] == "Stored report"
|
|
assert detail["evidence"]["sources"] == ([] if oversized else [source])
|
|
assert detail["evidence"]["omitted_count"] == 2 + int(oversized)
|
|
assert "prompt" not in detail
|
|
assert "execution_spec" not in detail
|
|
|
|
|
|
def test_browser_evidence_counts_valid_captured_excerpt_filtered_at_read():
|
|
from langchain_core.messages import ToolMessage
|
|
|
|
from deerflow.community.ragflow.formatting import format_retrieval_sources
|
|
from deerflow.community.ragflow.sources import durable_source_artifact
|
|
from deerflow.subagents.step_events import capture_step_message
|
|
|
|
content, artifact = format_retrieval_sources(
|
|
{"chunks": [{"id": "chunk", "dataset_id": "kb", "document_id": "doc", "content": "x" * 17_000}]},
|
|
dataset_names_by_id={"kb": "Knowledge base"},
|
|
max_chars_per_chunk=20_000,
|
|
max_total_chars=30_000,
|
|
)
|
|
captured = []
|
|
capture_step_message(ToolMessage(content=content, artifact=artifact, name="knowledge_search", tool_call_id="search", id="message"), captured, set())
|
|
source_id = artifact["knowledge_sources"]["sources"][0]["id"]
|
|
persisted = durable_source_artifact(captured, f"Claim [citation:1](#knowledge-{source_id})", max_chars=100_000)
|
|
assert len(persisted["knowledge_sources"]["sources"][0]["text"]) == 17_000
|
|
assert persisted["knowledge_sources"]["omitted_count"] == 0
|
|
assert service_module._browser_evidence_snapshot(persisted) == {"version": 1, "sources": [], "omitted_count": 1}
|
|
|
|
|
|
@pytest.mark.parametrize("persisted_omissions", [0, 3])
|
|
def test_browser_evidence_source_cap_preserves_omissions(persisted_omissions):
|
|
source = {"id": "a" * 32 + "-1", "provider": "ragflow", "dataset_id": "kb", "document_id": "doc", "chunk_id": "chunk", "dataset_name": "KB", "document_name": "Manual", "text": "Excerpt", "pages": [], "truncated": False}
|
|
artifact = {"knowledge_sources": {"version": 1, "sources": [source] * 103, "omitted_count": persisted_omissions}}
|
|
projected = service_module._browser_evidence_snapshot(artifact)
|
|
assert len(projected["sources"]) == 100
|
|
assert projected["omitted_count"] == persisted_omissions + 3
|
|
assert artifact["knowledge_sources"]["omitted_count"] == persisted_omissions
|
|
|
|
|
|
@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")
|