From 882433ae503d2d64ac6469136bc505a4ae4672dc Mon Sep 17 00:00:00 2001 From: jswSpace <1140534808@qq.com> Date: Fri, 30 Oct 2026 21:50:39 +0800 Subject: [PATCH] fix(subagents): recover expired attempts in paused batches (#7686) * fix(subagents): recover expired attempts in paused batches * test(subagents): target item reads in lease recovery race --------- Co-authored-by: jswSpace <198684035+jswSpace@users.noreply.github.com> --- README.md | 2 + .../harness/deerflow/persistence/AGENTS.md | 8 + .../persistence/subagent_batches/sql.py | 26 ++- backend/tests/AGENTS.md | 7 + backend/tests/test_batch_lease_lifecycle.py | 19 +- .../tests/test_paused_batch_lease_recovery.py | 173 ++++++++++++++++++ 6 files changed, 230 insertions(+), 5 deletions(-) create mode 100644 backend/tests/test_paused_batch_lease_recovery.py diff --git a/README.md b/README.md index 0141a8d3ea..6a2ff036b5 100644 --- a/README.md +++ b/README.md @@ -2795,6 +2795,8 @@ Explicitly enabled nested `task` delegation propagates omitted opaque source IDs If an item's final lease expires after exhausting the retry budget, the batch reaches a terminal state once all items are terminal: `failed` when none succeeded, or `completed` when at least one succeeded. The item remains `failed` in both cases so partial results stay visible. +Paused batches also reconcile expired attempts after a worker exits. Retryable items remain queued without starting new work until you explicitly resume the batch; exhausted attempts become failed and can be retried individually. + Direct `create_deerflow_agent(...)` integrations can own the same boundary explicitly instead of relying on Gateway startup. Construct one `SubagentRuntime` and share it across every graph in that application; its `max_running`, ordinary per-run total, bound `task` tool, and optional durable-batch tools then use the same caller-owned snapshot and execution controller. A runtime with a batch repository owns a worker and must be started before graph construction and stopped during application shutdown: ```python diff --git a/backend/packages/harness/deerflow/persistence/AGENTS.md b/backend/packages/harness/deerflow/persistence/AGENTS.md index e83b230eca..a96fd657ab 100644 --- a/backend/packages/harness/deerflow/persistence/AGENTS.md +++ b/backend/packages/harness/deerflow/persistence/AGENTS.md @@ -12,6 +12,14 @@ Alembic stamp/upgrade workers started inside `bootstrap_schema()` remain owned b On SQLite, `BEGIN IMMEDIATE` takes the database-wide write lock, so every unrelated writer (run status, thread metadata, the scheduler) waits for the transaction and fails with `database is locked` after `busy_timeout` (30s). Do slow work such as document conversion before opening a locked transaction; lock only to revalidate and publish, as `ProjectDocumentRepository.publish_under_live_lock` does. `tests/test_project_document_tools.py::TestConversionSerialization` pins this. +## Batch lease recovery + +Batch lease recovery includes paused parents only when they have expired active +items. Reconcile through the existing owner/expiry fence, then skip promotion and +claiming while paused. Preserve explicit retry/resume and terminal projection; +`tests/test_paused_batch_lease_recovery.py` covers HTTP, restart and export, and +`test_batch_lease_lifecycle.py` covers a concurrent renewal. + ## Unique-key upserts A repository method that creates or updates a row by a unique key must survive two first writers: both can miss the SELECT, and the later INSERT raises `IntegrityError`, a 500 for the caller. Use a dialect `on_conflict_do_update` when every key column is non-null, or retry the select-then-write after rolling back the session, as `ChannelConnectionRepository.upsert_connection` / `set_thread_id` and `FeedbackRepository.upsert` do. Feedback keeps the retry because `user_id` is nullable and a NULL key never conflicts, so ON CONFLICT would stop updating the NULL-user row. diff --git a/backend/packages/harness/deerflow/persistence/subagent_batches/sql.py b/backend/packages/harness/deerflow/persistence/subagent_batches/sql.py index f22ee29fe7..d847d6c1a7 100644 --- a/backend/packages/harness/deerflow/persistence/subagent_batches/sql.py +++ b/backend/packages/harness/deerflow/persistence/subagent_batches/sql.py @@ -5,7 +5,7 @@ from collections import Counter from datetime import UTC, datetime, timedelta from typing import Any -from sqlalchemy import func, select, update +from sqlalchemy import func, or_, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm import defer @@ -250,12 +250,27 @@ class SubagentBatchRepository: limit: int, exclude_item_ids: tuple[str, ...] = (), ) -> list[dict[str, Any]]: - """Promote pending work and atomically claim runnable items.""" + """Recover expired attempts, then promote and claim only runnable work.""" if limit <= 0: return [] claimed: list[dict[str, Any]] = [] async with self._sf() as session: - batches = list((await session.execute(select(SubagentBatchRow).where(SubagentBatchRow.status.in_(("queued", "running"))).order_by(SubagentBatchRow.created_at, SubagentBatchRow.id).with_for_update(skip_locked=True))).scalars()) + expired_items = select(SubagentBatchItemRow.id).where( + SubagentBatchItemRow.batch_id == SubagentBatchRow.id, + SubagentBatchItemRow.status.in_(("leased", "running")), + SubagentBatchItemRow.lease_expires_at < now, + SubagentBatchItemRow.id.not_in(exclude_item_ids), + ) + batches = list( + ( + await session.execute( + select(SubagentBatchRow) + .where(or_(SubagentBatchRow.status.in_(("queued", "running")), (SubagentBatchRow.status == "paused") & expired_items.exists())) + .order_by(SubagentBatchRow.created_at, SubagentBatchRow.id) + .with_for_update(skip_locked=True) + ) + ).scalars() + ) for batch in batches: if len(claimed) >= limit: break @@ -305,6 +320,11 @@ class SubagentBatchRepository: if terminalized_expired: await self._refresh_batch_status(session, batch, now=now) + # Pausing blocks new work, not recovery of an abandoned attempt. + # Retryable items remain queued until an explicit resume. + if batch.status == "paused": + continue + counts = await self._counts(session, batch.id) live = counts["queued"] + counts["leased"] + counts["running"] promote_count = max(0, batch.max_live_items - live) diff --git a/backend/tests/AGENTS.md b/backend/tests/AGENTS.md index c5095e9731..7e88e5dd75 100644 --- a/backend/tests/AGENTS.md +++ b/backend/tests/AGENTS.md @@ -220,6 +220,13 @@ contract probes execute the actual migration test setup through SQLAlchemy's dialect argument conversion, stopping before database acquisition; the real SQLite/PostgreSQL cases still exercise upgrade, downgrade and re-upgrade. +## Batch lease renewal races + +`test_batch_lease_lifecycle.py` parks recovery after the per-item expiry SELECT, +identified by selected item columns; the parent query can contain the same expiry +predicate. Assert one rejected item UPDATE as well as the preserved row, so an +empty expiry scan cannot satisfy the concurrency guard. + ## DingTalk stream tasks DingTalk stream lifecycle tests drive `_serve_stream` with a fake websocket and diff --git a/backend/tests/test_batch_lease_lifecycle.py b/backend/tests/test_batch_lease_lifecycle.py index f507d43627..0f67226d1a 100644 --- a/backend/tests/test_batch_lease_lifecycle.py +++ b/backend/tests/test_batch_lease_lifecycle.py @@ -308,17 +308,31 @@ async def test_each_scheduler_claim_has_a_fresh_ownership_generation(repository, @pytest.mark.asyncio -async def test_recovery_does_not_overwrite_a_renewal_after_its_expiry_read(repository, monkeypatch): +@pytest.mark.parametrize("batch_paused", [False, True]) +async def test_recovery_does_not_overwrite_a_renewal_after_its_expiry_read(repository, monkeypatch, batch_paused): await create_batch(repository) now = datetime.now(UTC) item = (await repository.claim_items(now=now, lease_owner="live-worker", lease_seconds=10, limit=1))[0] + if batch_paused: + await repository.pause_batch("batch-lease-test", user_id="user-1") paused, release = asyncio.Event(), asyncio.Event() original = AsyncSession.execute recovery = None + recovery_update_rowcounts = [] async def execute(session, statement, *args, **kwargs): result = await original(session, statement, *args, **kwargs) - if asyncio.current_task() is recovery and not paused.is_set() and getattr(statement, "is_select", False) and "lease_expires_at <" in str(statement): + if asyncio.current_task() is recovery and getattr(statement, "is_update", False) and statement.table.name == SubagentBatchItemRow.__tablename__: + recovery_update_rowcounts.append(result.rowcount) + if ( + asyncio.current_task() is recovery + and not paused.is_set() + and getattr(statement, "is_select", False) + and "lease_expires_at <" in str(statement) + and any(getattr(column, "table", None) is SubagentBatchItemRow.__table__ for column in statement.selected_columns) + ): + # The parent SELECT also contains the expiry predicate via EXISTS. + # Park only after the expired item itself has been read. paused.set() await release.wait() return result @@ -332,6 +346,7 @@ async def test_recovery_does_not_overwrite_a_renewal_after_its_expiry_read(repos finally: release.set() claimed = await asyncio.wait_for(recovery, 3) + assert recovery_update_rowcounts == [0], "recovery must attempt an item update and reject the renewed lease" assert claimed == [] assert await row_snapshot(repository, item["id"]) == expected diff --git a/backend/tests/test_paused_batch_lease_recovery.py b/backend/tests/test_paused_batch_lease_recovery.py new file mode 100644 index 0000000000..347a612235 --- /dev/null +++ b/backend/tests/test_paused_batch_lease_recovery.py @@ -0,0 +1,173 @@ +"""Paused batches still reconcile abandoned attempts without starting new work.""" + +import json +from datetime import UTC, datetime, timedelta +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import httpx +import pytest +import pytest_asyncio +from _router_auth_helpers import make_authed_test_app +from fastapi import FastAPI + +from app.gateway.auth_disabled import AUTH_SOURCE_SESSION +from app.gateway.routers import subagent_batches +from deerflow.config.database_config import DatabaseConfig +from deerflow.config.subagent_batches_config import SubagentBatchesConfig +from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig +from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config +from deerflow.persistence.subagent_batches import SubagentBatchRepository +from deerflow.persistence.subagent_batches.model import SubagentBatchItemRow +from deerflow.persistence.thread_meta import ThreadMetaRepository +from deerflow.subagents.batch_service import SubagentBatchService + +_URL = "/api/threads/thread-1/subagent-batches/batch-1" + + +async def _open_database(app: FastAPI, tmp_path: Path) -> tuple[SubagentBatchRepository, SubagentBatchService]: + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + sessions = get_session_factory() + app.state.subagent_batch_repo = SubagentBatchRepository(sessions) + app.state.thread_store = ThreadMetaRepository(sessions) + service = SubagentBatchService( + repository=app.state.subagent_batch_repo, + config=SubagentBatchesConfig(lease_seconds=10), + runtime_config=SubagentRuntimeConfig(), + ) + # No executor is needed: a paused batch must never dispatch its items. + service._execute_item = AsyncMock() + app.state.subagent_batch_service = service + return app.state.subagent_batch_repo, service + + +@pytest_asyncio.fixture +async def api(tmp_path): + # Synthetic session identity; routing, owner checks and SQL are real. + user = SimpleNamespace(id="user-1", system_role="user") + app = make_authed_test_app(user_factory=lambda: user, bind_current_user=True) + app.include_router(subagent_batches.router) + + @app.middleware("http") + async def session_source(request, call_next): + request.state.auth_source = AUTH_SOURCE_SESSION + return await call_next(request) + + try: + await _open_database(app, tmp_path) + await app.state.thread_store.create("thread-1", user_id="user-1") + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + yield app, client, tmp_path + finally: + await close_engine() + + +async def _create(repo: SubagentBatchRepository, *, count: int, max_attempts: int) -> None: + await repo.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="Paused research", + subagent_type="general-purpose", + items=[{"key": f"item-{index}", "prompt": f"Process {index}"} for index in range(count)], + max_live_items=1, + max_running_items=1, + max_attempts=max_attempts, + execution_spec={}, + ) + + +async def _finish(repo: SubagentBatchRepository, item: dict, *, owner: str, now: datetime) -> None: + assert await repo.finalize_item( + item["id"], + lease_owner=owner, + succeeded=True, + result="Saved report", + result_preview="Saved report", + result_truncated=False, + error=None, + stop_reason=None, + token_usage=None, + model_name=None, + completed_at=now, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("max_attempts", [1, 2]) +async def test_paused_batch_recovers_expired_item_without_dispatch_after_restart(api, max_attempts): + app, client, tmp_path = api + repo = app.state.subagent_batch_repo + await _create(repo, count=2, max_attempts=max_attempts) + now = datetime.now(UTC) + lost = (await repo.claim_items(now=now, lease_owner="lost-worker", lease_seconds=10, limit=1))[0] + assert await repo.mark_item_running(lost["id"], lease_owner="lost-worker", now=now) + assert (await client.post(f"{_URL}/pause")).json()["status"] == "paused" + + # Reopen the same file, with no heartbeat/finalization from the lost worker. + await close_engine() + repo, service = await _open_database(app, tmp_path) + await service.run_once(now=now + timedelta(seconds=10)) + before = (await client.get(f"{_URL}/items")).json() + assert before[0]["status"] == "running" # Equality is still a live lease. + + await service.run_once(now=now + timedelta(seconds=11)) + response = await client.get(f"{_URL}/items") + assert response.status_code == 200 + items = response.json() + assert [item["status"] for item in items] == ["failed" if max_attempts == 1 else "queued", "pending"] + assert items[0]["attempt"] == 1 + async with get_session_factory()() as session: + stored = await session.get(SubagentBatchItemRow, lost["id"]) + assert stored.lease_owner is None + assert stored.lease_expires_at is None + assert (await client.get(_URL)).json()["status"] == "paused" + service._execute_item.assert_not_awaited() + + if max_attempts == 1: + retried = await client.post(f"{_URL}/items/{lost['id']}/retry") + assert retried.status_code == 200 + assert retried.json()["attempt"] == 0 + await service.run_once(now=now + timedelta(seconds=12)) + assert (await client.get(_URL)).json()["status"] == "paused" + service._execute_item.assert_not_awaited() + + # Explicit resume is still required to claim retryable work. + assert (await client.post(f"{_URL}/resume")).json()["status"] == "queued" + resumed = (await repo.claim_items(now=now + timedelta(seconds=13), lease_owner="new-worker", lease_seconds=10, limit=1))[0] + assert resumed["id"] == lost["id"] + assert resumed["attempt"] == (1 if max_attempts == 1 else 2) + await _finish(repo, resumed, owner="new-worker", now=now + timedelta(seconds=14)) + exported = await client.get(f"{_URL}/results.jsonl") + rows = [json.loads(line) for line in exported.text.splitlines()] + assert rows[0]["result"] == "Saved report" + assert rows[1]["status"] == "pending" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("live_completion", [False, True]) +@pytest.mark.parametrize("succeeded_sibling", [False, True]) +async def test_last_paused_item_reaches_terminal_status_on_completion_or_expiry(api, live_completion, succeeded_sibling): + app, client, _ = api + repo = app.state.subagent_batch_repo + await _create(repo, count=2 if succeeded_sibling else 1, max_attempts=1) + now = datetime.now(UTC) + item = (await repo.claim_items(now=now, lease_owner="worker", lease_seconds=10, limit=1))[0] + if succeeded_sibling: + await _finish(repo, item, owner="worker", now=now) + item = (await repo.claim_items(now=now, lease_owner="worker", lease_seconds=10, limit=1))[0] + assert (await client.post(f"{_URL}/pause")).json()["status"] == "paused" + if live_completion: + await _finish(repo, item, owner="worker", now=now + timedelta(seconds=1)) + else: + await app.state.subagent_batch_service.run_once(now=now + timedelta(seconds=11)) + + batch = (await client.get(_URL)).json() + assert batch["status"] == ("completed" if live_completion or succeeded_sibling else "failed") + assert batch["completed_at"] is not None + assert batch["counts"]["leased"] == batch["counts"]["running"] == 0 + app.state.subagent_batch_service._execute_item.assert_not_awaited()