From 3db2fc8bb49a43be5a39fafc31f9d6c197973a46 Mon Sep 17 00:00:00 2001 From: yi-coding Date: Thu, 15 Oct 2026 12:02:00 +0800 Subject: [PATCH] fix(subagents): fence durable batch lease lifecycles (#6764) * fix(subagents): fence durable batch lease lifecycles * fix(subagents): fence batch claims and require captured lease owners --------- Co-authored-by: Willem Jiang --- backend/README.md | 9 + .../persistence/subagent_batches/sql.py | 288 ++++++------ .../harness/deerflow/subagents/AGENTS.md | 9 + .../harness/deerflow/subagents/batch_lease.py | 65 +++ .../deerflow/subagents/batch_service.py | 79 ++-- .../blocking_io/test_tool_assembly_offloop.py | 2 +- .../tests/test_batch_extension_snapshot.py | 4 +- backend/tests/test_batch_lease_lifecycle.py | 422 ++++++++++++++++++ backend/tests/test_subagent_batch_service.py | 10 +- 9 files changed, 690 insertions(+), 198 deletions(-) create mode 100644 backend/packages/harness/deerflow/subagents/batch_lease.py create mode 100644 backend/tests/test_batch_lease_lifecycle.py diff --git a/backend/README.md b/backend/README.md index 59b59cac83..56e8917eb2 100644 --- a/backend/README.md +++ b/backend/README.md @@ -94,6 +94,15 @@ Async task delegation with concurrent execution: - **Execution**: Background thread pools with status tracking and SSE events - **Flow**: Agent calls `task()` tool → executor runs subagent in background → polls for completion → returns result +Durable `batch_task` items keep their lease during tool initialization, execution, +and acceptance checks. Each claim uses a fresh ownership token; a recovered +attempt cannot be overwritten by an earlier worker's renewal or result. Workers +exclude their own active items from local recovery, while other workers can still +recover abandoned leases. Expiry recovery rechecks the captured lease so it cannot +overwrite a heartbeat that renewed after the expiry read. Pending promotion and +queued claims use conditional database writes, so competing workers cannot claim +the same item or undo a pause/cancellation observed after their initial read. + ### Memory System LLM-powered persistent context retention across conversations: diff --git a/backend/packages/harness/deerflow/persistence/subagent_batches/sql.py b/backend/packages/harness/deerflow/persistence/subagent_batches/sql.py index 8d31bb53ec..f22ee29fe7 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 +from sqlalchemy import func, select, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from sqlalchemy.orm import defer @@ -248,6 +248,7 @@ class SubagentBatchRepository: lease_owner: str, lease_seconds: int, limit: int, + exclude_item_ids: tuple[str, ...] = (), ) -> list[dict[str, Any]]: """Promote pending work and atomically claim runnable items.""" if limit <= 0: @@ -258,6 +259,7 @@ class SubagentBatchRepository: for batch in batches: if len(claimed) >= limit: break + runnable_batch = select(SubagentBatchRow.id).where(SubagentBatchRow.id == batch.id, SubagentBatchRow.status.in_(("queued", "running"))) expired = list( ( @@ -267,6 +269,7 @@ class SubagentBatchRepository: SubagentBatchItemRow.batch_id == batch.id, SubagentBatchItemRow.status.in_(("leased", "running")), SubagentBatchItemRow.lease_expires_at < now, + SubagentBatchItemRow.id.not_in(exclude_item_ids), ) .with_for_update(skip_locked=True) ) @@ -274,21 +277,30 @@ class SubagentBatchRepository: ) terminalized_expired = False for item in expired: - item.lease_owner = None - item.lease_expires_at = None - item.updated_at = now + values = {"lease_owner": None, "lease_expires_at": None, "updated_at": now} + terminal = item.cancel_requested_at is not None or item.attempt >= batch.max_attempts if item.cancel_requested_at is not None: - item.status = "cancelled" - item.completed_at = now - terminalized_expired = True + values.update(status="cancelled", completed_at=now) elif item.attempt >= batch.max_attempts: - item.status = "failed" - item.error = item.error or "Execution lease expired after the maximum retry count" - item.completed_at = now - terminalized_expired = True + values.update(status="failed", error=item.error or "Execution lease expired after the maximum retry count", completed_at=now) else: - item.status = "queued" - item.error = "Previous worker lease expired; retrying" + values.update(status="queued", error="Previous worker lease expired; retrying") + recovered = await session.execute( + update(SubagentBatchItemRow) + .where( + SubagentBatchItemRow.id == item.id, + SubagentBatchItemRow.status.in_(("leased", "running")), + SubagentBatchItemRow.lease_owner == item.lease_owner, + SubagentBatchItemRow.lease_expires_at == item.lease_expires_at, + ) + .values(**values) + .execution_options(synchronize_session=False) + ) + if recovered.rowcount == 1: + # A heartbeat may have renewed after SQLite's unlocked + # expiry read. Only the unchanged expired claim is ours. + await session.refresh(item) + terminalized_expired |= terminal if terminalized_expired: await self._refresh_batch_status(session, batch, now=now) @@ -312,8 +324,19 @@ class SubagentBatchRepository: ).scalars() ) for item in pending: - item.status = "queued" - item.updated_at = now + promoted = await session.execute( + update(SubagentBatchItemRow) + .where( + SubagentBatchItemRow.id == item.id, + SubagentBatchItemRow.status == "pending", + SubagentBatchItemRow.cancel_requested_at.is_(None), + SubagentBatchItemRow.batch_id.in_(runnable_batch), + ) + .values(status="queued", updated_at=now) + .execution_options(synchronize_session=False) + ) + if promoted.rowcount == 1: + await session.refresh(item) counts = await self._counts(session, batch.id) batch_available = max(0, batch.max_running_items - counts["leased"] - counts["running"]) @@ -328,6 +351,7 @@ class SubagentBatchRepository: SubagentBatchItemRow.batch_id == batch.id, SubagentBatchItemRow.status == "queued", SubagentBatchItemRow.cancel_requested_at.is_(None), + SubagentBatchItemRow.id.not_in(exclude_item_ids), ) .order_by(SubagentBatchItemRow.position) .limit(take) @@ -336,74 +360,96 @@ class SubagentBatchRepository: ).scalars() ) expires_at = now + timedelta(seconds=lease_seconds) + batch_claimed = False for item in runnable: - item.status = "leased" - item.attempt += 1 - item.lease_owner = lease_owner - item.lease_expires_at = expires_at - item.started_at = now - item.updated_at = now - item.error = None + # SQLite's SELECT FOR UPDATE is unlocked: only a successful + # conditional write grants this pass ownership of the item. + won = await session.execute( + update(SubagentBatchItemRow) + .where( + SubagentBatchItemRow.id == item.id, + SubagentBatchItemRow.status == "queued", + SubagentBatchItemRow.cancel_requested_at.is_(None), + SubagentBatchItemRow.batch_id.in_(runnable_batch), + ) + .values(status="leased", attempt=SubagentBatchItemRow.attempt + 1, lease_owner=lease_owner, lease_expires_at=expires_at, started_at=now, updated_at=now, error=None) + .execution_options(synchronize_session=False) + ) + if won.rowcount != 1: + continue + await session.refresh(item) value = self._item_dict(item) value["prompt"] = item.prompt value["batch"] = self._execution_batch_dict(batch) claimed.append(value) - if runnable: + batch_claimed = True + if batch_claimed: + await session.refresh(batch) batch.status = "running" batch.updated_at = now await session.commit() return claimed - async def renew_item_lease( - self, - item_id: str, - *, - lease_owner: str, - lease_seconds: int, - now: datetime, - ) -> dict[str, bool]: - async with self._sf() as session: - item = ( - await session.execute( - select(SubagentBatchItemRow) - .where( - SubagentBatchItemRow.id == item_id, - SubagentBatchItemRow.status.in_(("leased", "running")), - SubagentBatchItemRow.lease_owner == lease_owner, - ) - .with_for_update() + async def _owned_item_for_update(self, session: AsyncSession, item_id: str, lease_owner: str, *, statuses: tuple[str, ...] = ("leased", "running")) -> tuple[SubagentBatchItemRow | None, SubagentBatchRow | None]: + # Match claim/control paths' parent-before-item lock order on PostgreSQL. + batch_id = await session.scalar(select(SubagentBatchItemRow.batch_id).where(SubagentBatchItemRow.id == item_id)) + if batch_id is None: + return None, None + batch = await session.get(SubagentBatchRow, batch_id, with_for_update=True) + item = ( + await session.execute( + select(SubagentBatchItemRow) + .where( + SubagentBatchItemRow.id == item_id, + SubagentBatchItemRow.lease_owner == lease_owner, + SubagentBatchItemRow.status.in_(statuses), ) - ).scalar_one_or_none() - if item is None: + .with_for_update() + ) + ).scalar_one_or_none() + return item, batch + + @staticmethod + async def _write_owned_item(session: AsyncSession, item_id: str, lease_owner: str, values: dict[str, Any], *, statuses: tuple[str, ...] = ("leased", "running")) -> bool: + # SQLite ignores FOR UPDATE. Fence at the write itself, not only the + # earlier SELECT; never flush a dirty ORM snapshot over a newer claim. + result = await session.execute( + update(SubagentBatchItemRow) + .where( + SubagentBatchItemRow.id == item_id, + SubagentBatchItemRow.lease_owner == lease_owner, + SubagentBatchItemRow.status.in_(statuses), + ) + .values(**values) + .execution_options(synchronize_session=False) + ) + return result.rowcount == 1 + + async def renew_item_lease(self, item_id: str, *, lease_owner: str, lease_seconds: int, now: datetime) -> dict[str, bool]: + async with self._sf() as session: + item, batch = await self._owned_item_for_update(session, item_id, lease_owner) + if item is None or batch is None: + cancel_requested = batch is not None and (batch.status == "cancelled" or await session.scalar(select(SubagentBatchItemRow.cancel_requested_at).where(SubagentBatchItemRow.id == item_id)) is not None) + return {"valid": False, "cancel_requested": cancel_requested} + if item.cancel_requested_at is not None or batch.status == "cancelled": return {"valid": False, "cancel_requested": True} - batch = await session.get(SubagentBatchRow, item.batch_id) - cancel_requested = item.cancel_requested_at is not None or batch is None or batch.status == "cancelled" - if not cancel_requested: - item.lease_expires_at = now + timedelta(seconds=lease_seconds) - item.updated_at = now - await session.commit() - return {"valid": not cancel_requested, "cancel_requested": cancel_requested} + valid = await self._write_owned_item(session, item_id, lease_owner, {"lease_expires_at": now + timedelta(seconds=lease_seconds), "updated_at": now}) + cancel_requested = False + if not valid: + await session.refresh(item) + await session.refresh(batch) + cancel_requested = item.cancel_requested_at is not None or batch.status == "cancelled" + await session.commit() + return {"valid": valid, "cancel_requested": cancel_requested} async def mark_item_running(self, item_id: str, *, lease_owner: str, now: datetime) -> bool: async with self._sf() as session: - item = ( - await session.execute( - select(SubagentBatchItemRow) - .where( - SubagentBatchItemRow.id == item_id, - SubagentBatchItemRow.status == "leased", - SubagentBatchItemRow.lease_owner == lease_owner, - ) - .with_for_update() - ) - ).scalar_one_or_none() + item, _batch = await self._owned_item_for_update(session, item_id, lease_owner, statuses=("leased",)) if item is None or item.cancel_requested_at is not None: return False - item.status = "running" - item.started_at = now - item.updated_at = now + valid = await self._write_owned_item(session, item_id, lease_owner, {"status": "running", "started_at": now, "updated_at": now}, statuses=("leased",)) await session.commit() - return True + return valid async def finalize_item( self, @@ -423,99 +469,63 @@ class SubagentBatchRepository: result_artifact: dict[str, Any] | None = None, ) -> bool: async with self._sf() as session: - item = ( - await session.execute( - select(SubagentBatchItemRow) - .where( - SubagentBatchItemRow.id == item_id, - SubagentBatchItemRow.status.in_(("leased", "running")), - SubagentBatchItemRow.lease_owner == lease_owner, - ) - .with_for_update() - ) - ).scalar_one_or_none() + item, batch = await self._owned_item_for_update(session, item_id, lease_owner) if item is None: return False - batch = await session.get(SubagentBatchRow, item.batch_id, with_for_update=True) cancelled = item.cancel_requested_at is not None or batch is None or batch.status == "cancelled" - item.lease_owner = None - item.lease_expires_at = None - item.model_name = model_name - item.stop_reason = stop_reason - item.token_usage = token_usage - item.updated_at = completed_at - item.acceptance_verdict = None - item.result_artifact = None + values = { + "lease_owner": None, + "lease_expires_at": None, + "model_name": model_name, + "stop_reason": stop_reason, + "token_usage": token_usage, + "updated_at": completed_at, + "acceptance_verdict": None, + "result_artifact": None, + } if cancelled: - item.status = "cancelled" - item.error = "Cancelled by user" - item.completed_at = completed_at + values.update(status="cancelled", error="Cancelled by user", completed_at=completed_at) elif succeeded: - item.status = "succeeded" - item.acceptance_verdict = validate_acceptance_verdict(acceptance_verdict) - item.result = result - item.result_artifact = result_artifact - item.result_preview = result_preview - item.result_truncated = result_truncated - item.error = None - item.completed_at = completed_at + values.update( + status="succeeded", + acceptance_verdict=validate_acceptance_verdict(acceptance_verdict), + result=result, + result_artifact=result_artifact, + result_preview=result_preview, + result_truncated=result_truncated, + error=None, + completed_at=completed_at, + ) elif item.attempt < batch.max_attempts: - item.status = "queued" - item.error = error - item.started_at = None + values.update(status="queued", error=error, started_at=None) else: - item.status = "failed" - item.error = error - item.completed_at = completed_at + values.update(status="failed", error=error, completed_at=completed_at) + if not await self._write_owned_item(session, item_id, lease_owner, values): + return False if batch is not None: + # A SQLite parent read has no row lock: preserve a concurrent + # pause before deriving the batch projection under our write lock. + await session.refresh(batch) await self._refresh_batch_status(session, batch, now=completed_at) await session.commit() return True - async def requeue_item_after_admission_failure( - self, - item_id: str, - *, - lease_owner: str, - error: str | None, - now: datetime, - ) -> bool: - """Undo a claim rejected before execution admission. - - Claiming increments ``attempt`` so crash recovery can bound real - executions. A process-wide capacity rejection happens before an - execution starts, so it must release the lease and restore that - attempt instead of consuming the batch's retry budget. - """ + async def requeue_item_after_admission_failure(self, item_id: str, *, lease_owner: str, error: str | None, now: datetime) -> bool: + """Release unadmitted work without consuming its execution attempt.""" async with self._sf() as session: - item = ( - await session.execute( - select(SubagentBatchItemRow) - .where( - SubagentBatchItemRow.id == item_id, - SubagentBatchItemRow.status.in_(("leased", "running")), - SubagentBatchItemRow.lease_owner == lease_owner, - ) - .with_for_update() - ) - ).scalar_one_or_none() + item, batch = await self._owned_item_for_update(session, item_id, lease_owner) if item is None: return False - batch = await session.get(SubagentBatchRow, item.batch_id, with_for_update=True) cancelled = item.cancel_requested_at is not None or batch is None or batch.status == "cancelled" - item.lease_owner = None - item.lease_expires_at = None - item.updated_at = now + values = {"lease_owner": None, "lease_expires_at": None, "updated_at": now} if cancelled: - item.status = "cancelled" - item.error = "Cancelled by user" - item.completed_at = now + values.update(status="cancelled", error="Cancelled by user", completed_at=now) else: - item.status = "queued" - item.attempt = max(0, item.attempt - 1) - item.started_at = None - item.error = error + values.update(status="queued", attempt=max(0, item.attempt - 1), started_at=None, error=error) + if not await self._write_owned_item(session, item_id, lease_owner, values): + return False if batch is not None: + await session.refresh(batch) await self._refresh_batch_status(session, batch, now=now) await session.commit() return True diff --git a/backend/packages/harness/deerflow/subagents/AGENTS.md b/backend/packages/harness/deerflow/subagents/AGENTS.md index 4df3ff98dd..3c664d03fb 100644 --- a/backend/packages/harness/deerflow/subagents/AGENTS.md +++ b/backend/packages/harness/deerflow/subagents/AGENTS.md @@ -59,6 +59,15 @@ verdicts through queries and JSONL exports: absent criteria, checker errors, legacy rows, and failed executions are never implicitly accepted. Acceptance does not change retry policy. +**Batch lease supervision**: `batch_lease.py` owns one drained heartbeat across tool +assembly, execution and acceptance. Capture a fresh bounded claim owner each pass; +exclude locally active IDs from recovery. Execution requires that explicit token. +Promotion/claim writes fence item and parent status; worker writes fence owner and +active status in SQL. Renewal reports cancellation only for explicit cancel state. +Expiry recovery compares expiry. Keep PostgreSQL locks parent-first and polling bounded +by lease/3. `tests/test_batch_lease_lifecycle.py` covers real SQLite interleavings, +threaded initialization, peer cancellation and renewal drain. + **Built-in Agents**: `general-purpose` (all tools except `task`) and `bash` (command specialist) **Registry and managed definitions**: Runtime resolution is built-in → `config.yaml custom_agents` → enabled administrator-managed definitions, followed by explicit `subagents.agents.` overrides. Managed definitions are deployment-wide, persist through the same `agent_storage.backend` selection as Custom Agent definitions, and remain stored but are excluded from runtime when a built-in or later-added config definition owns the same name. The default Lead Agent sees the whole enabled catalog. A Custom Agent's `allowed_subagents` is snapshotted into run metadata (`None` = all, `[]` = hard deny, list = allowlist) and must filter both prompt discovery and `task` execution; never reload caller policy from mutable agent config inside the tool. **Benefit-based routing policy**: Direct execution is the default. Permit `task` when parallel latency, specialist skill, or context isolation clearly outweighs startup, duplicate discovery, synthesis, state conflicts, and side effects. Parallel scopes need independent outputs and disjoint mutable state; use the fewest useful agents and reassess each batch. A bounded sequential subagent is allowed for strong specialist/context benefit. With a per-response limit of 1, prompts omit parallel and multi-batch benefits and allow only material specialist/context benefit. Keep `lead_agent/prompt.py`, the `task` description, and built-in roles aligned; see `tests/test_subagent_routing_prompt.py`, `tests/test_subagent_prompt_security.py`, and `tests/test_lead_agent_prompt.py`. diff --git a/backend/packages/harness/deerflow/subagents/batch_lease.py b/backend/packages/harness/deerflow/subagents/batch_lease.py new file mode 100644 index 0000000000..e6688b2b6b --- /dev/null +++ b/backend/packages/harness/deerflow/subagents/batch_lease.py @@ -0,0 +1,65 @@ +"""Owned lease heartbeat for one durable batch execution.""" + +from __future__ import annotations + +import asyncio +import logging +from datetime import UTC, datetime + +logger = logging.getLogger(__name__) + + +class BatchItemLease: + """Keep initialization and execution leased until their supervisor exits.""" + + def __init__(self, repository, *, item_id: str, owner: str, seconds: int) -> None: + self._repository = repository + self.item_id = item_id + self.owner = owner + self._seconds = seconds + self.valid = True + self._stop = asyncio.Event() + self._renew_lock = asyncio.Lock() + self._task: asyncio.Task[None] | None = None + + async def renew(self) -> bool: + async with self._renew_lock: + if not self.valid: + return False + try: + result = await self._repository.renew_item_lease(self.item_id, lease_owner=self.owner, lease_seconds=self._seconds, now=datetime.now(UTC)) + self.valid = result["valid"] + except Exception: + self.valid = False + logger.exception("Durable batch lease renewal failed (item_id=%s)", self.item_id) + return self.valid + + async def _heartbeat(self) -> None: + while not self._stop.is_set(): + try: + await asyncio.wait_for(self._stop.wait(), timeout=max(1.0, self._seconds / 3)) + except TimeoutError: + if not await self.renew(): + return + + async def __aenter__(self) -> BatchItemLease: + if await self.renew(): + self._task = asyncio.create_task(self._heartbeat(), name=f"batch-item-lease-{self.item_id}") + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + self._stop.set() + if self._task is None: + return + cancellation: asyncio.CancelledError | None = None + # Never leave a renewal writing after this execution releases ownership. + # Database calls retain their normal connection/statement timeouts. + while not self._task.done(): + try: + await asyncio.wait({self._task}) + except asyncio.CancelledError as error: + if cancellation is None: + cancellation = error + self._task.result() + if cancellation is not None: + raise cancellation diff --git a/backend/packages/harness/deerflow/subagents/batch_service.py b/backend/packages/harness/deerflow/subagents/batch_service.py index c6807edfc9..3e09250ead 100644 --- a/backend/packages/harness/deerflow/subagents/batch_service.py +++ b/backend/packages/harness/deerflow/subagents/batch_service.py @@ -15,6 +15,7 @@ from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig from deerflow.extensions import LoadedExtensions, get_loaded_extensions from deerflow.mcp_scope import THREAD_INCARNATION_CONTEXT_KEY from deerflow.subagents.batch_acceptance import check_batch_acceptance +from deerflow.subagents.batch_lease import BatchItemLease from deerflow.subagents.batch_runtime import BatchSubmitRequest from deerflow.subagents.capacity import SubagentExecutionCapacity from deerflow.subagents.config import SubagentConfig, resolve_subagent_model_name @@ -191,18 +192,21 @@ class SubagentBatchService: available = max(0, self._runtime_config.max_running - len(self._executions)) if available <= 0: return + # A process identity is not a fence for a later claim by that process. + claim_owner = f"{self._lease_owner[:80]}:{uuid.uuid4().hex}" items = await self._repository.claim_items( now=now, - lease_owner=self._lease_owner, + lease_owner=claim_owner, lease_seconds=self._config.lease_seconds, limit=available, + exclude_item_ids=tuple(self._executions), ) for item in items: item_id = item["id"] if item_id in self._executions: continue task = asyncio.create_task( - self._execute_item(item), + self._execute_item(item, lease_owner=claim_owner), name=f"subagent-batch-item-{item_id}", ) self._executions[item_id] = task @@ -280,7 +284,12 @@ class SubagentBatchService: # what lets another worker own the HTTP control request safely. return batch - async def _execute_item(self, item: dict[str, Any]) -> None: + async def _execute_item(self, item: dict[str, Any], *, lease_owner: str) -> None: + async with BatchItemLease(self._repository, item_id=item["id"], owner=lease_owner, seconds=self._config.lease_seconds) as lease: + if lease.valid: + await self._execute_leased_item(item, lease) + + async def _execute_leased_item(self, item: dict[str, Any], lease: BatchItemLease) -> None: item_id = item["id"] execution_id: str | None = None try: @@ -318,13 +327,7 @@ class SubagentBatchService: # assembly blocked in the worker thread — the poll loop's checks # only start after execute_async(), so launching without this # check would run work the user already cancelled. - lease = await self._repository.renew_item_lease( - item_id, - lease_owner=self._lease_owner, - lease_seconds=self._config.lease_seconds, - now=datetime.now(UTC), - ) - if not lease["valid"]: + if not await lease.renew(): logger.info( "Durable batch item %s cancelled or lease lost during tool assembly; skipping launch", item_id, @@ -357,13 +360,7 @@ class SubagentBatchService: execution_id = executor.execute_async(prompt, task_id=item_id) self._execution_ids[item_id] = execution_id marked_running = False - renew_every = max(1.0, self._config.lease_seconds / 3) - status_poll_every = min( - self._config.poll_interval_seconds, - renew_every, - ) - loop = asyncio.get_running_loop() - next_renew_at = loop.time() + renew_every + status_poll_every = min(self._config.poll_interval_seconds, max(1.0, self._config.lease_seconds / 3)) while True: result = get_background_task_result(execution_id) if result is None: @@ -371,40 +368,29 @@ class SubagentBatchService: if result.status is SubagentStatus.RUNNING and not marked_running: marked_running = await self._repository.mark_item_running( item_id, - lease_owner=self._lease_owner, + lease_owner=lease.owner, now=datetime.now(UTC), ) if not marked_running: request_cancel_background_task(execution_id) if result.status.is_terminal: break - now_monotonic = loop.time() - if now_monotonic >= next_renew_at: - lease = await self._repository.renew_item_lease( - item_id, - lease_owner=self._lease_owner, - lease_seconds=self._config.lease_seconds, - now=datetime.now(UTC), - ) - next_renew_at = loop.time() + renew_every - if not lease["valid"]: - request_cancel_background_task(execution_id) + if not lease.valid: + request_cancel_background_task(execution_id) try: - until_renew = max(0.0, next_renew_at - loop.time()) - await asyncio.wait_for( - self._stop.wait(), - timeout=min(status_poll_every, until_renew), - ) + await asyncio.wait_for(self._stop.wait(), timeout=status_poll_every) if self._stop.is_set(): raise asyncio.CancelledError except TimeoutError: pass + if not lease.valid: + return raw_result = result.result or "" if getattr(result, "admission_failure", False): await self._repository.requeue_item_after_admission_failure( item_id, - lease_owner=self._lease_owner, + lease_owner=lease.owner, error=result.error, now=datetime.now(UTC), ) @@ -420,7 +406,7 @@ class SubagentBatchService: acceptance_verdict = None if result.status is SubagentStatus.COMPLETED and item.get("acceptance_criteria"): try: - valid, acceptance_verdict = await self._check_acceptance_with_lease(item, result, app_config) + valid, acceptance_verdict = await self._check_acceptance_with_lease(item, result, app_config, lease) if not valid: return except Exception: @@ -429,7 +415,7 @@ class SubagentBatchService: logger.warning("Batch acceptance check failed; result remains unchecked (item_id=%s)", item_id, exc_info=True) await self._repository.finalize_item( item_id, - lease_owner=self._lease_owner, + lease_owner=lease.owner, succeeded=result.status is SubagentStatus.COMPLETED, result=stored_result, result_preview=preview, @@ -453,7 +439,7 @@ class SubagentBatchService: ) await self._repository.finalize_item( item_id, - lease_owner=self._lease_owner, + lease_owner=lease.owner, succeeded=False, result=None, result_preview=None, @@ -500,19 +486,10 @@ class SubagentBatchService: pending.add_done_callback(_pending_releases.discard) pending.add_done_callback(lambda done: _log_release_failure(done, execution_id=execution_id)) - async def _check_acceptance_with_lease(self, item, result, app_config): + async def _check_acceptance_with_lease(self, item, result, app_config, lease: BatchItemLease): """Keep a completed execution leased until its advisory check drains.""" - async def renew(): - lease = await self._repository.renew_item_lease( - item["id"], - lease_owner=self._lease_owner, - lease_seconds=self._config.lease_seconds, - now=datetime.now(UTC), - ) - return lease["valid"] - - if not await renew(): + if not await lease.renew(): return False, None check = asyncio.create_task( check_batch_acceptance( @@ -526,8 +503,8 @@ class SubagentBatchService: while True: done, _ = await asyncio.wait({check}, timeout=max(1.0, self._config.lease_seconds / 3)) if done: - return True, check.result() - if not await renew(): + return lease.valid, check.result() + if not lease.valid: return False, None finally: if not check.done(): diff --git a/backend/tests/blocking_io/test_tool_assembly_offloop.py b/backend/tests/blocking_io/test_tool_assembly_offloop.py index c683211a4b..c265833603 100644 --- a/backend/tests/blocking_io/test_tool_assembly_offloop.py +++ b/backend/tests/blocking_io/test_tool_assembly_offloop.py @@ -235,7 +235,7 @@ async def test_batch_item_assembles_off_loop(monkeypatch, tmp_path): }, } - await service._execute_item(item) + await service._execute_item(item, lease_owner="assembly-claim") assert len(finalize_calls) == 1 assert finalize_calls[0]["item_id"] == "item-1" diff --git a/backend/tests/test_batch_extension_snapshot.py b/backend/tests/test_batch_extension_snapshot.py index 9990acbad8..b1924fd942 100644 --- a/backend/tests/test_batch_extension_snapshot.py +++ b/backend/tests/test_batch_extension_snapshot.py @@ -93,7 +93,7 @@ async def test_batch_keeps_its_construction_snapshot(monkeypatch, _subagent_env, monkeypatch.setattr(_subagent_env.SubagentExecutor, "_create_agent", create) try: set_loaded_extensions(b if replacement else a) - await service._execute_item(item) + await service._execute_item(item, lease_owner="original-claim") assert repo.finalize_item.await_args.kwargs["succeeded"] if late_replace: set_loaded_extensions(b) @@ -107,7 +107,7 @@ async def test_batch_keeps_its_construction_snapshot(monkeypatch, _subagent_env, original_spec = deepcopy(item["batch"]["execution_spec"]) restarted = service_module.SubagentBatchService(repository=repo, config=SubagentBatchesConfig(), runtime_config=SubagentRuntimeConfig(), app_config=config, extensions=b) set_loaded_extensions(a) - await restarted._execute_item(item) + await restarted._execute_item(item, lease_owner="restarted-claim") executor = _subagent_env.SubagentExecutor(**captured["executor_kwargs"]) result = await executor._aexecute("identify") assert result.status == _subagent_env.SubagentStatus.COMPLETED, result.error diff --git a/backend/tests/test_batch_lease_lifecycle.py b/backend/tests/test_batch_lease_lifecycle.py new file mode 100644 index 0000000000..f507d43627 --- /dev/null +++ b/backend/tests/test_batch_lease_lifecycle.py @@ -0,0 +1,422 @@ +"""Real database regressions for durable batch lease generations and setup.""" + +import asyncio +import threading +from datetime import UTC, datetime, timedelta +from enum import Enum +from types import SimpleNamespace + +import pytest +import pytest_asyncio +from sqlalchemy.ext.asyncio import AsyncSession + +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, SubagentBatchRow +from deerflow.subagents import batch_service as service_module +from deerflow.subagents.batch_service import SubagentBatchService + + +class SubagentStatus(Enum): + RUNNING = "running" + COMPLETED = "completed" + + @property + def is_terminal(self): + return self is SubagentStatus.COMPLETED + + +@pytest_asyncio.fixture +async def repository(tmp_path): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + yield SubagentBatchRepository(get_session_factory()) + finally: + await close_engine() + + +async def create_batch(repository, *, count=1, max_attempts=3): + return await repository.create_batch( + batch_id="batch-lease-test", + user_id="user-1", + thread_id="thread-1", + run_id="run-1", + tool_call_id="call-1", + submission_key="run-1:call-1", + title="Lease recovery", + subagent_type="general-purpose", + items=[{"key": f"record-{i}", "prompt": f"Process {i}"} for i in range(count)], + max_live_items=count, + max_running_items=1, + max_attempts=max_attempts, + execution_spec={"subagent_config": {"name": "general-purpose", "description": "test", "system_prompt": "Work carefully."}}, + ) + + +async def row_snapshot(repository, item_id): + async with repository._sf() as session: + item = await session.get(SubagentBatchItemRow, item_id) + return {c.name: getattr(item, c.name) for c in item.__table__.columns} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["renew", "running", "finalize", "requeue"]) +@pytest.mark.parametrize("peer_action", ["reclaim", "cancel"]) +async def test_in_flight_old_acknowledgement_cannot_change_reclaimed_row(repository, monkeypatch, operation, peer_action): + await create_batch(repository) + now = datetime.now(UTC) + first = (await repository.claim_items(now=now, lease_owner="old-claim", lease_seconds=10, limit=1))[0] + item_id = first["id"] + paused, release = asyncio.Event(), asyncio.Event() + original = AsyncSession.execute + old_task = None + + async def execute(session, statement, *args, **kwargs): + result = await original(session, statement, *args, **kwargs) + if asyncio.current_task() is old_task and not paused.is_set() and getattr(statement, "is_select", False) and "subagent_batch_items" in str(statement) and len(statement.selected_columns) > 1: + # Pause after the old operation's actual database read, not before + # it starts. SQLite ignores SELECT FOR UPDATE; a peer can reclaim. + paused.set() + await release.wait() + return result + + monkeypatch.setattr(AsyncSession, "execute", execute) + + async def acknowledge(): + stamp = now + timedelta(seconds=1) # Within the OLD lease; expiry must not mask a missing fence. + if operation == "renew": + return await repository.renew_item_lease(item_id, lease_owner="old-claim", lease_seconds=10, now=stamp) + if operation == "running": + return await repository.mark_item_running(item_id, lease_owner="old-claim", now=stamp) + if operation == "requeue": + return await repository.requeue_item_after_admission_failure(item_id, lease_owner="old-claim", error="stale rejection", now=stamp) + return await repository.finalize_item( + item_id, + lease_owner="old-claim", + succeeded=True, + result="stale result", + result_preview="stale result", + result_truncated=False, + error=None, + stop_reason=None, + token_usage={"total_tokens": 99}, + model_name="old-model", + completed_at=stamp, + ) + + old_task = asyncio.create_task(acknowledge()) + try: + await asyncio.wait_for(paused.wait(), 3) + if peer_action == "reclaim": + replacement = await repository.claim_items(now=now + timedelta(seconds=11), lease_owner="new-claim", lease_seconds=10, limit=1) + assert replacement[0]["id"] == item_id and replacement[0]["attempt"] == 2 + else: + await repository.cancel_batch("batch-lease-test", user_id="user-1") + expected = await row_snapshot(repository, item_id) + finally: + release.set() + outcome = await asyncio.wait_for(old_task, 3) + assert outcome == ({"valid": False, "cancel_requested": peer_action == "cancel"} if operation == "renew" else False) + assert await row_snapshot(repository, item_id) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["pending", "queued"]) +@pytest.mark.parametrize("peer_action", ["claim", "cancel", "pause"]) +async def test_stale_claim_pass_preserves_peer_claim_or_control(repository, monkeypatch, phase, peer_action): + await create_batch(repository) + now = datetime.now(UTC) + item_id = (await repository.list_items("batch-lease-test", user_id="user-1"))[0]["id"] + if phase == "queued": + # Promote without claiming so the paused pass has not acquired SQLite's + # writer lock before its runnable-item read. + assert await repository.claim_items(now=now, lease_owner="seed", lease_seconds=10, limit=1, exclude_item_ids=(item_id,)) == [] + paused, release = asyncio.Event(), asyncio.Event() + original = AsyncSession.execute + stale = None + + async def execute(session, statement, *args, **kwargs): + result = await original(session, statement, *args, **kwargs) + if asyncio.current_task() is stale and not paused.is_set() and getattr(statement, "is_select", False) and "subagent_batch_items.status =" in str(statement) and len(statement.selected_columns) > 1: + if phase in statement.compile().params.values(): + paused.set() + await release.wait() + return result + + monkeypatch.setattr(AsyncSession, "execute", execute) + stale = asyncio.create_task(repository.claim_items(now=now, lease_owner="stale-claim", lease_seconds=10, limit=1)) + try: + await asyncio.wait_for(paused.wait(), 3) + if peer_action == "claim": + peer = await repository.claim_items(now=now + timedelta(seconds=1), lease_owner="peer-claim", lease_seconds=10, limit=1) + assert len(peer) == 1 and peer[0]["attempt"] == 1 + elif peer_action == "cancel": + await repository.cancel_batch("batch-lease-test", user_id="user-1") + else: + await repository.pause_batch("batch-lease-test", user_id="user-1") + expected = await row_snapshot(repository, item_id) + expected_batch = await repository.get_batch("batch-lease-test", user_id="user-1") + finally: + release.set() + claimed = await asyncio.wait_for(stale, 3) + assert claimed == [], "a stale read must not grant ownership over the peer's claim or control" + assert await row_snapshot(repository, item_id) == expected + assert await repository.get_batch("batch-lease-test", user_id="user-1") == expected_batch + + +@pytest.mark.asyncio +@pytest.mark.parametrize("case", ["missing", "wrong_owner", "item_cancelled", "batch_cancelled", "owned"]) +async def test_renewal_reports_only_explicit_cancellation(repository, case): + await create_batch(repository) + now = datetime.now(UTC) + item = (await repository.claim_items(now=now, lease_owner="claim", lease_seconds=10, limit=1))[0] + if case in ("item_cancelled", "batch_cancelled"): + async with repository._sf() as session: + if case == "item_cancelled": + (await session.get(SubagentBatchItemRow, item["id"])).cancel_requested_at = now + else: + (await session.get(SubagentBatchRow, "batch-lease-test")).status = "cancelled" + await session.commit() + outcome = await repository.renew_item_lease("missing" if case == "missing" else item["id"], lease_owner="peer" if case == "wrong_owner" else "claim", lease_seconds=10, now=now) + assert outcome == {"valid": case == "owned", "cancel_requested": case in ("item_cancelled", "batch_cancelled")} + + +@pytest.mark.asyncio +async def test_execution_requires_an_explicit_claim_owner(repository): + service = make_service(repository) + with pytest.raises(TypeError, match="lease_owner"): + await service._execute_item({"id": "missing"}) + + +def setup_executor(monkeypatch, started, release): + launched = [] + monkeypatch.setattr(service_module, "SubagentStatus", SubagentStatus) + + def assemble(**_kwargs): + started.set() + assert release.wait(10), "test must release the real assembly worker" + return [] + + class Executor: + def __init__(self, **_kwargs): + pass + + def execute_async(self, _prompt, task_id=None): + launched.append(task_id) + return "execution-1" + + result = SimpleNamespace(status=SubagentStatus.COMPLETED, result="correct source result", error=None, stop_reason=None, token_usage_records=None, ai_messages=[]) + monkeypatch.setattr("deerflow.tools.get_available_tools", assemble) + monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_a, **_kw: "test-model") + monkeypatch.setattr(service_module, "SubagentExecutor", Executor) + monkeypatch.setattr(service_module, "get_background_task_result", lambda _id: result) + monkeypatch.setattr(service_module, "cleanup_background_task", lambda _id: None) + return launched + + +def make_service(repository): + return SubagentBatchService(repository=repository, config=SubagentBatchesConfig(lease_seconds=10), runtime_config=SubagentRuntimeConfig(max_running=3), app_config=SimpleNamespace()) + + +@pytest.mark.asyncio +async def test_scheduler_does_not_reclaim_its_still_initializing_item(repository, monkeypatch): + await create_batch(repository, max_attempts=1) + started, release = threading.Event(), threading.Event() + launched = setup_executor(monkeypatch, started, release) + service = make_service(repository) + now = datetime.now(UTC) + await service.run_once(now=now) + tasks = list(service._executions.values()) + try: + assert await asyncio.to_thread(started.wait, 3) + # Advance the scheduler's logical claim clock past its lease, while a + # real assembly-pool worker still owns initialization. No private lock + # is held or manufactured by the test. + await service.run_once(now=now + timedelta(seconds=11)) + item = (await repository.list_items("batch-lease-test", user_id="user-1"))[0] + assert item["status"] == "leased", "live initialization must not exhaust its retry budget" + assert item["attempt"] == 1 + finally: + release.set() + await asyncio.wait_for(asyncio.gather(*tasks), 3) + assert len(launched) == 1 + item = (await repository.list_items("batch-lease-test", user_id="user-1", include_result=True))[0] + assert item["status"] == "succeeded" and item["result"] == "correct source result" + + +@pytest.mark.asyncio +async def test_initialization_keeps_renewing_for_other_workers(repository, monkeypatch): + await create_batch(repository) + started, release = threading.Event(), threading.Event() + setup_executor(monkeypatch, started, release) + service = make_service(repository) + renewals, twice = [], asyncio.Event() + original = repository.renew_item_lease + + async def renew(*args, **kwargs): + result = await original(*args, **kwargs) + renewals.append(kwargs["now"]) + if len(renewals) >= 2: + twice.set() + return result + + monkeypatch.setattr(repository, "renew_item_lease", renew) + now = datetime.now(UTC) + await service.run_once(now=now) + tasks = list(service._executions.values()) + try: + assert await asyncio.to_thread(started.wait, 3) + try: + await asyncio.wait_for(twice.wait(), 5) + except TimeoutError: + pytest.fail("a blocked tool assembly received no periodic lease heartbeat") + # Just beyond the first renewal's expiry, but before the second one's. + peer_time = renewals[0] + timedelta(seconds=10, milliseconds=1) + assert await repository.claim_items(now=peer_time, lease_owner="peer", lease_seconds=10, limit=1) == [] + finally: + release.set() + await asyncio.wait_for(asyncio.gather(*tasks), 3) + + +@pytest.mark.asyncio +async def test_each_scheduler_claim_has_a_fresh_ownership_generation(repository, monkeypatch): + await create_batch(repository, count=2) + started, release = threading.Event(), threading.Event() + release.set() + setup_executor(monkeypatch, started, release) + service = make_service(repository) + owners = [] + original = repository.claim_items + + async def claim(**kwargs): + owners.append(kwargs["lease_owner"]) + return await original(**kwargs) + + monkeypatch.setattr(repository, "claim_items", claim) + for _ in range(2): + await service.run_once(now=datetime.now(UTC)) + await asyncio.gather(*list(service._executions.values())) + await asyncio.sleep(0) # Deliver the normal task-registry done callbacks. + assert len(owners) == len(set(owners)) == 2, "worker identity must not be reused as a claim-generation fence" + assert all(len(owner) <= 128 for owner in owners) + items = await repository.list_items("batch-lease-test", user_id="user-1", include_result=True) + assert [item["status"] for item in items] == ["succeeded", "succeeded"] + assert all("lease_owner" not in item for item in items) + + +@pytest.mark.asyncio +async def test_recovery_does_not_overwrite_a_renewal_after_its_expiry_read(repository, monkeypatch): + 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] + paused, release = asyncio.Event(), asyncio.Event() + original = AsyncSession.execute + recovery = None + + 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): + paused.set() + await release.wait() + return result + + monkeypatch.setattr(AsyncSession, "execute", execute) + recovery = asyncio.create_task(repository.claim_items(now=now + timedelta(seconds=11), lease_owner="recoverer", lease_seconds=10, limit=1)) + try: + await asyncio.wait_for(paused.wait(), 3) + assert (await repository.renew_item_lease(item["id"], lease_owner="live-worker", lease_seconds=10, now=now + timedelta(seconds=12)))["valid"] + expected = await row_snapshot(repository, item["id"]) + finally: + release.set() + claimed = await asyncio.wait_for(recovery, 3) + assert claimed == [] + assert await row_snapshot(repository, item["id"]) == expected + + +@pytest.mark.asyncio +async def test_lost_heartbeat_cancels_execution_before_a_long_poll_interval(repository, monkeypatch): + await create_batch(repository) + started, release = threading.Event(), threading.Event() + release.set() + setup_executor(monkeypatch, started, release) + service = make_service(repository) + service._config = SubagentBatchesConfig(lease_seconds=10, poll_interval_seconds=60) + running, cancelled = asyncio.Event(), asyncio.Event() + mark_running = repository.mark_item_running + + async def mark(*args, **kwargs): + valid = await mark_running(*args, **kwargs) + if valid: + running.set() + return valid + + monkeypatch.setattr(repository, "mark_item_running", mark) + result = SimpleNamespace(status=SubagentStatus.RUNNING, result=None, error=None, stop_reason=None, token_usage_records=None, ai_messages=[]) + + def read_result(_id): + return result + + def cancel(_id): + result.status = SubagentStatus.COMPLETED + cancelled.set() + + monkeypatch.setattr(service_module, "get_background_task_result", read_result) + monkeypatch.setattr(service_module, "request_cancel_background_task", cancel) + await service.run_once(now=datetime.now(UTC)) + task = next(iter(service._executions.values())) + try: + await asyncio.wait_for(running.wait(), 3) + # Another worker handles the control request; no local cancellation shortcut. + await repository.cancel_batch("batch-lease-test", user_id="user-1") + await asyncio.wait_for(cancelled.wait(), 8) + await asyncio.wait_for(task, 4) + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + item = (await repository.list_items("batch-lease-test", user_id="user-1", include_result=True))[0] + assert item["status"] == "cancelled" and item["result"] is None + + +@pytest.mark.asyncio +async def test_heartbeat_exit_drains_an_in_flight_renewal_across_repeated_cancellation(): + from deerflow.subagents.batch_lease import BatchItemLease + + renewing, finish_renewal, exit_context = asyncio.Event(), asyncio.Event(), asyncio.Event() + calls = 0 + + class Repository: + async def renew_item_lease(self, *_args, **_kwargs): + nonlocal calls + calls += 1 + if calls > 1: + renewing.set() + await finish_renewal.wait() + return {"valid": True, "cancel_requested": False} + + async def supervise(): + async with BatchItemLease(Repository(), item_id="item", owner="generation", seconds=3): + await exit_context.wait() + + task = asyncio.create_task(supervise()) + try: + await asyncio.wait_for(renewing.wait(), 3) + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert not task.done(), "caller must retain its in-flight database renewal" + finish_renewal.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 3) + assert calls == 2 + assert not any(t.get_name() == "batch-item-lease-item" for t in asyncio.all_tasks()) + finally: + finish_renewal.set() + exit_context.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) diff --git a/backend/tests/test_subagent_batch_service.py b/backend/tests/test_subagent_batch_service.py index f89974bd28..95d172e99c 100644 --- a/backend/tests/test_subagent_batch_service.py +++ b/backend/tests/test_subagent_batch_service.py @@ -259,8 +259,8 @@ async def test_execute_item_polls_completion_without_waiting_for_lease_renewal(m raise AssertionError("a task that completes between polls need not expose running") async def renew_item_lease(self, *_args, **_kwargs): - # Exactly one pre-launch revalidation is expected; the poll loop - # must not renew for a task that completes between polls. + # 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} @@ -306,7 +306,7 @@ async def test_execute_item_polls_completion_without_waiting_for_lease_renewal(m assert repository.finalized is not None assert repository.finalized["result"] == "fast result" - assert repository.renews == 1 + assert repository.renews == 2 class _Registry: @@ -602,7 +602,7 @@ async def test_cancel_during_tool_assembly_skips_launch(monkeypatch, tmp_path) - ) claimed = await repository.claim_items( now=datetime.now(UTC), - lease_owner=service._lease_owner, + lease_owner="assembly-claim", lease_seconds=60, limit=10, ) @@ -648,7 +648,7 @@ async def test_cancel_during_tool_assembly_skips_launch(monkeypatch, tmp_path) - ), ) - item_task = asyncio.create_task(service._execute_item(item)) + 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")