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 <willem.jiang@gmail.com>

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:

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

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.<name>` 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`.

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

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():

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"

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

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)

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")