github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-10-26 13:16:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(runtime): attribute terminal commits and isolate lease renewals (#6919)

Browse Source 
* fix(runtime): confirm terminal commits and isolate lease renewal

* fix(runtime): diagnose unsupported terminal store contracts

* fix(runtime): attribute atomic claims and close lease lifecycle gaps

* fix(runtime): require explicit terminal completion capability

* fix(runtime): retry renewals that race rollback admission
This commit is contained in:
spud 2026-10-19 15:06:35 +08:00 committed by GitHub
parent 0a1c86a371
commit 850b3b2547
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
19 changed files with 2609 additions and 449 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
README.md
View File

@ -592,6 +592,22 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser
> [!IMPORTANT]
> The Gateway still owns active run tasks in process, so production defaults to a single Gateway worker (`GATEWAY_WORKERS=1`). Multi-worker deployments require Postgres, the Redis stream bridge (`stream_bridge.type: redis`), `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; process-local memory/JSONL event stores cannot enforce singleton delivery receipts across workers. Kubernetes replicas run one worker per Pod, which the worker count cannot see: declare them with `deployment.multi_instance: true` (or `DEER_FLOW_MULTI_INSTANCE=1`, which deploy tooling such as a Helm chart can set from its replica count) so the startup gate enforces the same prerequisites instead of staying inert. Declared instances that store credentials (`channel_connections`) must also share one at-rest `DEER_FLOW_CREDENTIALS_KEY` (the Helm chart and `make up` generate it; back it up, see [backend/docs/CONFIGURATION.md](backend/docs/CONFIGURATION.md#credentials-encryption-key-deer_flow_credentials_key)). The bridge shares SSE delivery and bounded `Last-Event-ID` replay across workers. When a valid reconnect cursor has been trimmed, or a subscriber that already established an empty-stream wait falls behind before its first delivery, Memory and Redis emit a machine-readable SSE `gap` event instead of silently returning a partial replay; the Web UI reloads durable thread/event state and resumes from the retained tail. An already-observed Redis stream removed during replay also signals a gap, so slow subscribers reload durable state instead of heartbeating forever after cleanup. Lease reconciliation marks runs from dead workers as errors, persists their delivery receipts, publishes the terminal stream marker, schedules retained-stream cleanup, and updates the affected thread status. SSE, `/wait`, and internal stream consumers use `stream_bridge.heartbeat_interval_seconds` (default `15`) for idle liveness checks; changing it requires a Gateway restart. Malformed Redis reconnect IDs live-tail new events instead of replaying the retained buffer, and the rolling retained-buffer TTL (`stream_ttl_seconds`) remains a cleanup safety net rather than a run timeout. Failed-login counters and lockouts for `POST /api/v1/auth/login/local` (`auth.local.max_login_attempts` / `lockout_seconds`) are kept in the shared `login_throttle` table whenever the application database is SQLite or Postgres (`auth.local.throttle_storage: auto`, the default), so every replica enforces one lockout per client IP; `memory` keeps the historical per-process counter, which under N replicas hands an attacker N × `max_login_attempts` guesses and logs a startup warning. IM chat-to-thread bindings live in the shared `channel_thread_bindings` table whenever the application database is SQLite or Postgres (an existing `channels/store.json` is imported once at startup and renamed `store.json.migrated`); IM bot credentials entered in the web UI live in the shared `channel_runtime_configs` table, encrypted with `DEER_FLOW_CREDENTIALS_KEY` (an existing `channels/runtime-config.json` is imported once and then deleted), so a save on any instance reaches the provider's connection within one lease renewal; each provider's platform connection is held by one instance through the shared `channel_leases` table — `channels.leader_lease` (default `auto`, on wherever the multi-process signal is) hands it to another instance when the holder dies, and the startup gate refuses `enabled: false` while an IM channel other than GitHub is enabled (see [IM channel leader lease](backend/docs/CONFIGURATION.md#im-channel-leader-lease-channelsleader_lease)); follow-up buffers still need their own multi-worker coordination. To try this topology on one machine, `make dev-multi` starts two Gateways on shared Postgres, Redis and `DEER_FLOW_HOME`, and `make dev-multi-check` runs the cross-instance checks; see the local two-Gateway harness in [backend/docs/CONFIGURATION.md](backend/docs/CONFIGURATION.md#local-two-gateway-harness).
>
> With run ownership heartbeats enabled, terminal success requires acknowledgment
> from that worker's own durable write. A matching database row alone cannot
> confirm success, and completion metadata cannot retry a rejected terminal
> transition. Each run renews independently, so a stalled store operation does
> not delay other runs. If acknowledgment remains unavailable when the last
> confirmed lease expires, the worker stops conservatively even if its commit
> actually reached the database. Expired leases cannot be renewed.
> Interrupt-and-replace admission acknowledges only the old runs actually claimed
> by its atomic transaction. Rollback keeps renewing its acknowledged interrupted
> row until restoration and the separate final error commit finish. An
> active-only renewal that raced that admission gets one bounded retry after the
> interrupted acknowledgment, without weakening the owner or live-lease checks.
> A failed renewal leaves a local expiry guard armed. Shutdown bounds its wait for orphan
> recovery, but cancellation-resistant storage may still be using its resources.
>
> The local two-Gateway harness disables IM providers under both `channels` and `channel_connections` by default. `DEERFLOW_MI_KEEP_CHANNELS=1` preserves them; `channels.leader_lease` stays as configured.
>

35
backend/packages/harness/deerflow/persistence/AGENTS.md
View File

@ -41,3 +41,38 @@ missing or strictly older leases remain recoverable.
## MCP task notifications
`mcp_tasks/sql.py::claim_notification_work` never returns a row without a dispatch event. A pending snapshot whose condition resolved before delivery (tracking recovered, or an answered `input_required` prompt returned the task to `working`) is consumed at claim time: `notified_version` advances, the content fingerprint clears so an identical later event is recorded again, stale error/attempt/dispatch state clears, status returns to `none`, and the claim lease is released instead of launching an empty event. A legacy row left without a dispatch event is consumed only when nothing newer exists; a newer recorded event is rebuilt and claimed instead. Consumption is a conditional UPDATE fenced on the event version and phase the claim read, so a concurrent newer event (SQLite has no row locks) is never acknowledged and stays claimable. Coverage: `tests/test_mcp_task_repository.py` (resolved snapshot, legacy retry, expired-claim reclaim, identical-event recurrence, interleaved newer completion, legacy launching/dispatched with and without a newer completion).
## Run terminal authority and lease expiry
`RunRepository.update_status(..., owner_worker_id=...)` and
`finalize_if_not_cancelled(..., owner_worker_id=...)` require that owner and a
live lease. Increment the change clock before taking the run row lock to retain
the existing clock/row lock order. PostgreSQL checks `clock_timestamp()` after
`SELECT FOR UPDATE`, rather than transaction-start `now()`; SQLite evaluates its
database clock in the conditional UPDATE after acquiring its writer lock.
`renew_lease` and legacy `update_lease` also reject expired existing leases and
already-expired target leases. `allow_interrupted=True` additionally permits
interrupted rows for the owning Worker's unfinished rollback; Manager requires
its own interrupted ACK and a rollback barrier before requesting it. Success/error
remain ineligible, and owner/live-lease conditions are identical. Atomic admission
returns immutable outcome snapshots of its committed claimed rows; a subsequent
interrupted -> error write must not mutate an earlier ACK (including in Memory).
Success returns only after transaction commit.
`update_run_completion(require_terminal=True, owner_worker_id=...)` requires the
exact terminal status/owner atomically and leaves outcome fields unchanged;
it cannot transition an active row. Metadata may arrive after the terminal lease
expires: it authorizes no execution or outcome transition, and Manager already
requires its own matching ACK. The owner/status predicate protects against a
changed row, but cannot distinguish deletion/recreation with the same run ID,
owner and outcome without a persistent generation. Third-party heartbeat stores
must implement these keyword contracts and atomic predicates, returning explicit
commit results; accepting `**kwargs` alone does not establish that guarantee.
For strict completion, explicitly declare both `require_terminal` and
`owner_worker_id`, or set `supports_terminal_completion_kwargs=True` on the class
defining `update_run_completion`. A flag-only child cannot opt in an inherited
method, and an override must declare its own opt-in. Manager skips unsupported
completion calls with a diagnostic before invoking the Store.
Unsupported owned calls fail closed; single-worker legacy calls remain supported.
These are scoped guards, not a generation
protocol covering every mutation. Contract/lock-wait tests live in
`tests/test_terminal_store_contract.py` (Postgres uses `TEST_POSTGRES_URI`).

77
backend/packages/harness/deerflow/persistence/run/sql.py
View File

@ -12,7 +12,7 @@ from datetime import UTC, datetime, timedelta
from typing import Any
from uuid import uuid4
from sqlalchemy import and_, case, delete, func, or_, select, text, update
from sqlalchemy import and_, case, delete, func, literal, or_, select, text, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
@ -63,6 +63,19 @@ class RunRepository(RunStore):
raise RuntimeError("run-change clock did not return a position")
return int(value)
@staticmethod
def _live_lease_predicate(session: AsyncSession, lease=RunRow.lease_expires_at):
# PostgreSQL clock_timestamp is evaluated after the row lock below,
# unlike transaction-start now(). SQLite evaluates its clock at UPDATE.
if session.bind.dialect.name == "postgresql":
return lease > func.clock_timestamp()
return func.julianday(lease) > func.julianday("now")
@staticmethod
async def _lock_owned_row(session: AsyncSession, run_id: str) -> None:
if session.bind.dialect.name == "postgresql":
await session.execute(select(RunRow.run_id).where(RunRow.run_id == run_id).with_for_update())
@staticmethod
def _normalize_model_name(model_name: str | None) -> str | None:
"""Normalize model_name for storage: strip whitespace, truncate to 128 chars."""
@ -376,7 +389,7 @@ class RunRepository(RunStore):
result = await session.execute(stmt)
return {row.run_id: self._row_to_dict(row) for row in result.scalars()}
async def update_status(self, run_id, status, *, error=None, stop_reason=None, goal_verdict=None) -> bool:
async def update_status(self, run_id, status, *, error=None, stop_reason=None, goal_verdict=None, owner_worker_id=None) -> bool:
values: dict[str, Any] = {"status": status, "updated_at": datetime.now(UTC)}
if error is not None:
values["error"] = error
@ -391,7 +404,15 @@ class RunRepository(RunStore):
# completed run) cannot be overwritten by a late writer.
async with self._sf() as session:
values["change_seq"] = await self._next_change_seq(session)
result = await session.execute(update(RunRow).where(RunRow.run_id == run_id, RunRow.status.in_(("pending", "running", "interrupted"))).values(**values))
conditions = [RunRow.run_id == run_id, RunRow.status.in_(("pending", "running", "interrupted"))]
if owner_worker_id is not None:
await self._lock_owned_row(session, run_id)
conditions.extend([RunRow.owner_worker_id == owner_worker_id, self._live_lease_predicate(session)])
if status != "error":
conditions.append(RunRow.status.in_(("pending", "running")))
if status == "success":
conditions.append(RunRow.cancel_action.is_(None))
result = await session.execute(update(RunRow).where(*conditions).values(**values))
await session.commit()
return result.rowcount != 0
@ -530,6 +551,8 @@ class RunRepository(RunStore):
run_id: str,
*,
status: str,
require_terminal: bool = False,
owner_worker_id: str | None = None,
total_input_tokens: int = 0,
total_output_tokens: int = 0,
total_tokens: int = 0,
@ -566,22 +589,28 @@ class RunRepository(RunStore):
values["last_ai_message"] = last_ai_message[:2000]
if first_human_message is not None:
values["first_human_message"] = first_human_message[:2000]
if error is not None:
if not require_terminal and error is not None:
values["error"] = error
if goal_verdict is not None:
if not require_terminal and goal_verdict is not None:
values["goal_verdict"] = self._safe_json(goal_verdict)
allowed_sources = ["pending", "running"]
if status not in allowed_sources:
allowed_sources.append(status)
if status == "error" and "interrupted" not in allowed_sources:
allowed_sources.append("interrupted")
if require_terminal and status in ("pending", "running"):
return False
async with self._sf() as session:
values["change_seq"] = await self._next_change_seq(session)
conditions = [RunRow.run_id == run_id]
if require_terminal:
conditions.extend([RunRow.status == status, RunRow.owner_worker_id == owner_worker_id])
else:
conditions.append(RunRow.status.in_(tuple(allowed_sources)))
result = await session.execute(
update(RunRow)
.where(
RunRow.run_id == run_id,
RunRow.status.in_(tuple(allowed_sources)),
*conditions,
)
.values(**values)
)
@ -721,6 +750,7 @@ class RunRepository(RunStore):
*,
owner_worker_id: str,
lease_expires_at: str,
allow_interrupted: bool = False,
) -> bool:
lease_dt = datetime.fromisoformat(lease_expires_at)
values: dict[str, Any] = {
@ -729,7 +759,18 @@ class RunRepository(RunStore):
"updated_at": datetime.now(UTC),
}
async with self._sf() as session:
result = await session.execute(update(RunRow).where(RunRow.run_id == run_id, RunRow.owner_worker_id == owner_worker_id, RunRow.status.in_(("pending", "running"))).values(**values))
await self._lock_owned_row(session, run_id)
result = await session.execute(
update(RunRow)
.where(
RunRow.run_id == run_id,
RunRow.owner_worker_id == owner_worker_id,
RunRow.status.in_(("pending", "running", "interrupted") if allow_interrupted else ("pending", "running")),
self._live_lease_predicate(session),
self._live_lease_predicate(session, literal(lease_dt)),
)
.values(**values)
)
await session.commit()
return result.rowcount != 0
@ -739,16 +780,20 @@ class RunRepository(RunStore):
*,
owner_worker_id: str,
lease_expires_at: str,
allow_interrupted: bool = False,
) -> LeaseRenewal:
"""Renew the owner lease and read cancellation intent atomically."""
lease_dt = datetime.fromisoformat(lease_expires_at)
async with self._sf() as session:
await self._lock_owned_row(session, run_id)
result = await session.execute(
update(RunRow)
.where(
RunRow.run_id == run_id,
RunRow.owner_worker_id == owner_worker_id,
RunRow.status.in_(("pending", "running")),
RunRow.status.in_(("pending", "running", "interrupted") if allow_interrupted else ("pending", "running")),
self._live_lease_predicate(session),
self._live_lease_predicate(session, literal(lease_dt)),
)
.values(
lease_expires_at=lease_dt,
@ -801,6 +846,7 @@ class RunRepository(RunStore):
error: str | None = None,
stop_reason: str | None = None,
goal_verdict: dict[str, Any] | None = None,
owner_worker_id: str | None = None,
) -> StatusFinalization:
"""Atomically let completion win only before cancellation."""
values: dict[str, Any] = {
@ -816,12 +862,14 @@ class RunRepository(RunStore):
async with self._sf() as session:
values["change_seq"] = await self._next_change_seq(session)
conditions = [RunRow.run_id == run_id, RunRow.status.in_(("pending", "running")), RunRow.cancel_action.is_(None)]
if owner_worker_id is not None:
await self._lock_owned_row(session, run_id)
conditions.extend([RunRow.owner_worker_id == owner_worker_id, self._live_lease_predicate(session)])
result = await session.execute(
update(RunRow)
.where(
RunRow.run_id == run_id,
RunRow.status.in_(("pending", "running")),
RunRow.cancel_action.is_(None),
*conditions,
)
.values(**values)
.returning(RunRow.run_id)
@ -830,7 +878,10 @@ class RunRepository(RunStore):
await session.commit()
return StatusFinalization(finalized=True)
current = await session.execute(select(RunRow.cancel_action).where(RunRow.run_id == run_id))
cancellation_conditions = [RunRow.run_id == run_id]
if owner_worker_id is not None:
cancellation_conditions.extend([RunRow.owner_worker_id == owner_worker_id, RunRow.status.in_(("pending", "running")), self._live_lease_predicate(session)])
current = await session.execute(select(RunRow.cancel_action).where(*cancellation_conditions))
cancel_action = current.scalar_one_or_none()
await session.commit()
return StatusFinalization(

46
backend/packages/harness/deerflow/runtime/AGENTS.md
View File

@ -102,9 +102,49 @@ answers for paginated history. See `docs/skill-usage-ui.md`.
**Deferred terminal commit:** With an event store, the worker stages its terminal
status locally and commits it only after finalization's receipt and duration
writes. `RunRecord.terminal_commit_pending` keeps `_renew_leases()` renewing that
still-active row until the commit is attempted; a renewal rejected by the worker's
own commit is confirmed by re-reading the row, while a peer claim fences the run.
Never select runs for renewal by local status alone. Shutdown's drain treats the
still-active row until its own terminal write is acknowledged or ownership is
lost. Manager-owned attempts register separate futures before Store I/O and
publish an immutable outcome proof only on explicit successful commit results;
matching row fields and legacy `None` results cannot confirm a commit. Atomic
replacement registers attempts before admission I/O and acknowledges only the
returned claimed rows, using their committed outcome fields. It does not repeat
the interrupted update. An interrupted ACK is intermediate while a live Worker
still owes rollback: keep a separate rollback barrier and explicitly renew that
interrupted row under the same live-owner/lease predicates until the final error
ACK. An active-only renewal that was already in flight across atomic rollback
admission is retried once with `allow_interrupted=True` only after this Worker's
interrupted ACK is available, and only while that same per-run decision remains
inside its original deadline. An in-flight admission ACK is awaited; row state is
never reread as attribution. Peer takeover, an expired lease, a changed/final
outcome, or a second rejection still fences. Even without an event store, defer
rollback's error write until restoration has finished. `tests/test_stale_active_renewal.py`
covers the Manager transition (including an ACK still in flight), and
`tests/test_atomic_admission_worker.py` covers the real Worker restoration/error
tail with and without an event Store. A fenced Worker must not begin a rollback checkpoint restoration,
including when fencing itself delivered cancellation. Admission compensation never retries an unowned write in heartbeat
mode. A local deadline guard remains armed after fast renewal errors and after
heartbeat stop; it rechecks the latest lease/proof and fences without waiting for
the admission lock. Record removal and final ACK cancel its timer. Bounded orphan
recovery drain requests cancellation, retains/observes unfinished I/O, and returns
without joining a cancellation-resistant Store. Neither that return nor heartbeat
stop authorizes closing still-in-use resources. A rejected
CAS never falls back to another terminal writer. In heartbeat mode completion
updates require that proof and an atomic matching terminal status/owner in Store;
they only enrich metadata. Single-worker completion retains its recovery fallback.
Never select runs for renewal by local status alone. Each run has a single-flight
renewal check and supervised I/O, bounded by its last confirmed durable lease and
a monotonic wait budget. Heartbeat ticks do not join slow checks; cancellation
signals are delivered per run as results arrive. Timeout requests cancellation
without joining cancellation-resistant I/O, whose late results are observed but
cannot restore local authority. Shutdown drains within its shared budget and
reports outstanding I/O rather than claiming it stopped. SQL/Memory reject expired
renewals. `stop_heartbeat()` is not a Store-resource idle barrier: a driver that
suppresses cancellation can retain its connection until the supervised I/O settles
or the process exits; closing that resource cannot be claimed safe from this return.
Owned terminal writes require a live owner lease (Postgres checks DB time
after its row lock). These guards do not fence every RunStore mutation or provide
a persistent attempt identity. Permanently missing ACKs may conservatively fence
a committed run. Shutdown's drain treats the
same state as active, and also awaits committed terminal workers still executing
hooks, END or resource cleanup. The registered worker shields and drains one owned
finalizer across repeated host cancellation, preserving the first interruption even

797
backend/packages/harness/deerflow/runtime/runs/manager.py
View File

@ -209,6 +209,39 @@ def _store_accepts_idempotency_parameter(store: RunStore, parameter_name: str) -
return bool(implementation_owner and implementation_owner.__dict__.get("supports_idempotency_request_kwargs", False))
def _store_supports_terminal_completion(store: RunStore) -> bool:
"""Require explicit strict keywords or an opt-in on the method's owner."""
callable_ = store.update_run_completion
if getattr(callable_, "__func__", None) is RunStore.update_run_completion:
return False
try:
parameters = inspect.signature(callable_).parameters
except (TypeError, ValueError):
return False
required = [parameters.get(name) for name in ("require_terminal", "owner_worker_id")]
if all(parameter is not None and parameter.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) for parameter in required):
return True
if any(parameter is not None and parameter.kind is inspect.Parameter.POSITIONAL_ONLY for parameter in required):
return False
if not any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()):
return False
implementation_owner = next(
(candidate for candidate in type(store).__mro__ if "update_run_completion" in candidate.__dict__),
None,
)
return bool(implementation_owner and implementation_owner.__dict__.get("supports_terminal_completion_kwargs", False))
@dataclass(frozen=True)
class TerminalCommitProof:
"""Process-local evidence returned by one completed durable write."""
status: RunStatus
error: str | None
stop_reason: str | None
goal_verdict: dict[str, Any] | None
@dataclass
class RunRecord:
"""Mutable record for a single run."""
@ -262,6 +295,10 @@ class RunRecord:
# memory and commits it only at the end of finalization, so the durable row
# stays active and its lease must keep renewing until that commit.
terminal_commit_pending: bool = False
# An interrupted ACK does not authorize completion of an unfinished rollback.
rollback_finalization_pending: bool = False
terminal_commit_proof: TerminalCommitProof | None = field(default=None, repr=False)
terminal_status_revision: int = field(default=0, repr=False)
idempotency_key: str | None = None
idempotency_request: dict[str, Any] | None = None
# True only on the caller that recovered an existing idempotent admission;
@ -319,6 +356,10 @@ class RunManager:
self._heartbeat_task: asyncio.Task | None = None
self._heartbeat_stop: asyncio.Event | None = None
self._orphan_recovery_task: asyncio.Task[None] | None = None
self._terminal_attempts: dict[str, set[asyncio.Future[bool]]] = {}
self._lease_checks: dict[str, asyncio.Task[None]] = {}
self._lease_operations: dict[str, asyncio.Task] = {}
self._lease_deadlines: dict[str, asyncio.TimerHandle] = {}
def _index_run_locked(self, record: RunRecord) -> None:
"""Register *record* in the thread index. Caller must hold ``self._lock``."""
@ -326,6 +367,7 @@ class RunManager:
def _unindex_run_locked(self, run_id: str, thread_id: str) -> None:
"""Drop *run_id* from the thread index. Caller must hold ``self._lock``."""
self._clear_lease_deadline(run_id)
bucket = self._runs_by_thread.get(thread_id)
if bucket is not None:
bucket.pop(run_id, None)
@ -449,6 +491,116 @@ class RunManager:
self._store_put_payload(record, error=error),
)
@staticmethod
def _terminal_outcome(record: RunRecord, status: RunStatus | None = None, *, error: str | None = None, stop_reason: str | None = None, goal_verdict: dict[str, Any] | None = None) -> TerminalCommitProof:
return TerminalCommitProof(
status or record.status,
error if error is not None else record.error,
stop_reason if stop_reason is not None else record.stop_reason,
copy.deepcopy(goal_verdict if goal_verdict is not None else record.goal_verdict),
)
def _has_terminal_ack(self, record: RunRecord, outcome: TerminalCommitProof | None = None) -> bool:
return not record.ownership_lost and (outcome or self._terminal_outcome(record)) == record.terminal_commit_proof
def _has_final_terminal_ack(self, record: RunRecord) -> bool:
return not record.rollback_finalization_pending and self._has_terminal_ack(record)
def can_release_terminal_barrier(self, record: RunRecord) -> bool:
"""Keep durable ownership barriers through rollback's final error ACK."""
return not self.heartbeat_enabled or self._store is None or record.ownership_lost or self._has_final_terminal_ack(record)
def _begin_terminal_attempt(self, record: RunRecord) -> asyncio.Future[bool]:
future = asyncio.get_running_loop().create_future()
self._terminal_attempts.setdefault(record.run_id, set()).add(future)
record.terminal_commit_pending = True
return future
def _end_terminal_attempt(self, record: RunRecord, future: asyncio.Future[bool], committed: bool) -> None:
future.set_result(committed)
attempts = self._terminal_attempts[record.run_id]
attempts.remove(future)
if not attempts:
self._terminal_attempts.pop(record.run_id)
def _acknowledge_terminal(self, record: RunRecord, outcome: TerminalCommitProof, staged_revision: int) -> None:
if self._runs.get(record.run_id) is not record or record.ownership_lost:
return
# Publish the acknowledged outcome without an await. An older
# interrupted ACK cannot supersede a newer staged/acknowledged error.
previous = record.terminal_commit_proof
if previous is None or previous.status == RunStatus.interrupted or outcome.status != RunStatus.interrupted:
record.terminal_commit_proof = outcome
authoritative = record.terminal_commit_proof
if authoritative.status != RunStatus.interrupted or record.terminal_status_revision == staged_revision:
record.status = authoritative.status
record.error = authoritative.error
record.stop_reason = authoritative.stop_reason
record.goal_verdict = copy.deepcopy(authoritative.goal_verdict)
if authoritative.status == RunStatus.error:
record.rollback_finalization_pending = False
record.terminal_status_persisted = self._has_final_terminal_ack(record)
if record.terminal_status_persisted:
record.terminal_commit_pending = False
self._clear_lease_deadline(record.run_id)
record.updated_at = _now_iso()
async def _terminal_operation(self, record: RunRecord, outcome: TerminalCommitProof, operation: Callable[[], Awaitable[Any]], *, cas: bool = False) -> Any:
future = self._begin_terminal_attempt(record)
staged_revision = record.terminal_status_revision
committed = False
try:
result = await operation()
committed = result.finalized is True if cas else result is True
if committed:
self._acknowledge_terminal(record, outcome, staged_revision)
return result
finally:
self._end_terminal_attempt(record, future, committed)
async def _atomic_replacement(self, candidates: list[RunRecord], action: str, operation: Callable[[], Awaitable[Any]]) -> Any:
"""Attribute returned claims while the caller holds the admission lock."""
attempts = [(r, r.terminal_commit_pending, self._begin_terminal_attempt(r), r.terminal_status_revision) for r in candidates]
claimed: dict[str, dict] = {}
try:
result = await operation()
claimed = {row["run_id"]: row for row in result[1]}
for record, _, _, revision in attempts:
row = claimed.get(record.run_id)
if row is None or row.get("status") != "interrupted" or row.get("owner_worker_id") != self._worker_id or row.get("thread_id") != record.thread_id:
continue
if record.ownership_lost:
continue
# The error writer may have ACKed while this older admission
# was awaiting its response. Do not restart cancellation after
# that final outcome, even though this claim also committed.
if record.terminal_commit_proof is not None and record.terminal_commit_proof.status != RunStatus.interrupted:
continue
task_active = record.task is not None and not record.task.done()
record.rollback_finalization_pending = action == "rollback" and task_active
outcome = TerminalCommitProof(RunStatus.interrupted, row.get("error"), row.get("stop_reason"), copy.deepcopy(row.get("goal_verdict")))
self._acknowledge_terminal(record, outcome, revision)
record.abort_action = action
record.abort_event.set()
if task_active and not record.finalizing:
record.task.cancel()
record.finalizing = task_active
return result
finally:
for record, was_pending, future, revision in attempts:
row = claimed.get(record.run_id)
committed = row is not None and row.get("status") == "interrupted" and row.get("owner_worker_id") == self._worker_id and row.get("thread_id") == record.thread_id
self._end_terminal_attempt(record, future, committed)
if not committed and not self._terminal_attempts.get(record.run_id) and record.terminal_status_revision == revision and not record.ownership_lost:
record.terminal_commit_pending = was_pending
async def _fence_unconfirmed_terminal(self, record: RunRecord) -> None:
deadline = self._parse_lease_deadline(record.lease_expires_at)
remaining = max(0.0, (deadline - datetime.now(UTC)).total_seconds()) if deadline else 0.0
if await self._release_confirmed_terminal_barrier(record, timeout=remaining):
return
await self._mark_ownership_lost(record, reason="This worker's terminal write was rejected or could not be confirmed.", require_active=False)
async def _persist_status(self, record: RunRecord, status: RunStatus, *, error: str | None = None, stop_reason: str | None = None) -> bool:
"""Best-effort persist a status transition to the backing store."""
if record.ownership_lost:
@ -460,6 +612,39 @@ class RunManager:
return False
if self._store is None:
return True
if self.heartbeat_enabled and not record.store_only:
terminal = status not in (RunStatus.pending, RunStatus.running)
outcome = self._terminal_outcome(record, status, error=error, stop_reason=stop_reason)
if terminal and self._has_terminal_ack(record, outcome):
return True
async def operation():
return await self._call_store_with_retry(
"update_status",
record.run_id,
lambda: self._store.update_status(
record.run_id,
status.value,
error=outcome.error,
stop_reason=outcome.stop_reason,
owner_worker_id=self._worker_id,
**({"goal_verdict": copy.deepcopy(outcome.goal_verdict)} if outcome.goal_verdict is not None else {}),
),
)
try:
updated = await self._terminal_operation(record, outcome, operation) if terminal else await operation()
if updated is True and not record.ownership_lost:
return True
except asyncio.CancelledError:
if terminal:
await self._fence_unconfirmed_terminal(record)
raise
except Exception:
logger.warning("Unable to confirm owned status write for run %s", record.run_id, exc_info=True)
if terminal:
await self._fence_unconfirmed_terminal(record)
return False
row_recovery_payload = self._store_put_payload(record, error=error, stop_reason=stop_reason)
try:
updated = await self._call_store_with_retry(
@ -567,6 +752,13 @@ class RunManager:
if record is not None and record.ownership_lost:
logger.warning("Skipped completion persistence for run %s after lease ownership was lost", run_id)
return
if self.heartbeat_enabled and self._store is not None:
if record is None or record.store_only or not self._has_final_terminal_ack(record) or kwargs.get("status") != record.status.value:
logger.warning("Skipped completion data without an attributable terminal ack for run %s", run_id)
return
if any(kwargs.get(key) is not None and kwargs[key] != getattr(record, key) for key in ("error", "stop_reason", "goal_verdict")):
logger.warning("Skipped completion data that changes the acknowledged outcome for run %s", run_id)
return
if record is not None:
if record.goal_verdict is not None:
# This fallback can be the first successful terminal write
@ -582,6 +774,28 @@ class RunManager:
row_recovery_payload = self._store_put_payload(record, error=kwargs.get("error"))
if self._store is None:
return
if self.heartbeat_enabled:
if record is None or record.store_only or not self._has_final_terminal_ack(record) or kwargs.get("status") != record.status.value:
logger.warning("Skipped completion data without an attributable terminal ack for run %s", run_id)
return
if not _store_supports_terminal_completion(self._store):
logger.warning(
"Skipped terminal completion metadata for run %s: Store %s must explicitly support require_terminal and owner_worker_id or declare supports_terminal_completion_kwargs on the class defining update_run_completion",
run_id,
type(self._store).__name__,
)
return
try:
updated = await self._call_store_with_retry(
"update_run_completion",
run_id,
lambda: self._store.update_run_completion(run_id, require_terminal=True, owner_worker_id=self._worker_id, **kwargs),
)
if updated is not True:
logger.warning("Terminal completion metadata was not updated for run %s", run_id)
except Exception:
logger.warning("Failed to persist terminal completion metadata for run %s", run_id, exc_info=True)
return
try:
updated = await self._call_store_with_retry(
"update_run_completion",
@ -1108,9 +1322,14 @@ class RunManager:
run_id,
)
return
proof = record.terminal_commit_proof
if self.heartbeat_enabled and proof is not None and proof.status != RunStatus.interrupted and self._terminal_outcome(record, status, error=error, stop_reason=stop_reason, goal_verdict=goal_verdict) != proof:
logger.info("Skipped status change after acknowledged terminal outcome for run %s", run_id)
return
record.status = status
record.updated_at = _now_iso()
if status not in (RunStatus.pending, RunStatus.running):
record.terminal_status_revision += 1
record.terminal_status_persisted = self._store is None
if error is not None:
record.error = error
@ -1122,7 +1341,7 @@ class RunManager:
persisted = await self._persist_status(record, status, error=error, stop_reason=stop_reason)
if status not in (RunStatus.pending, RunStatus.running):
async with self._lock:
record.terminal_status_persisted = persisted
record.terminal_status_persisted = self._has_final_terminal_ack(record) if self.heartbeat_enabled and self._store is not None else persisted
if not persisted and self.heartbeat_enabled and status == RunStatus.success and not record.ownership_lost:
await self._mark_ownership_lost(
record,
@ -1146,7 +1365,7 @@ class RunManager:
persisted = await self._persist_status(record, status, error=error, stop_reason=stop_reason)
if status not in (RunStatus.pending, RunStatus.running):
async with self._lock:
record.terminal_status_persisted = persisted
record.terminal_status_persisted = self._has_final_terminal_ack(record) if self.heartbeat_enabled and self._store is not None else persisted
if not persisted and self.heartbeat_enabled and status == RunStatus.success and not record.ownership_lost:
await self._mark_ownership_lost(
record,
@ -1177,52 +1396,52 @@ class RunManager:
)
return None
async with self._lock:
record = self._runs.get(run_id)
if record is None or record.store_only or record.ownership_lost:
return None
outcome = self._terminal_outcome(record, status, error=error, stop_reason=stop_reason, goal_verdict=goal_verdict)
if self._has_terminal_ack(record, outcome):
return None
if record.terminal_commit_proof is not None and record.terminal_commit_proof.status != RunStatus.interrupted:
return None
try:
result = await self._call_store_with_retry(
"finalize_if_not_cancelled",
run_id,
lambda: self._store.finalize_if_not_cancelled(
result = await self._terminal_operation(
record,
outcome,
lambda: self._call_store_with_retry(
"finalize_if_not_cancelled",
run_id,
status=status.value,
error=error,
stop_reason=stop_reason,
**({"goal_verdict": goal_verdict} if goal_verdict is not None else {}),
lambda: self._store.finalize_if_not_cancelled(
run_id,
status=status.value,
error=outcome.error,
stop_reason=outcome.stop_reason,
owner_worker_id=self._worker_id,
**({"goal_verdict": copy.deepcopy(outcome.goal_verdict)} if outcome.goal_verdict is not None else {}),
),
),
cas=True,
)
except Exception:
async with self._lock:
record = self._runs.get(run_id)
if record is not None:
record.terminal_status_persisted = False
if record is not None:
await self._mark_ownership_lost(
record,
reason=("The durable store could not confirm whether cancellation or completion won."),
require_active=False,
)
except asyncio.CancelledError:
await self._fence_unconfirmed_terminal(record)
raise
except Exception as exc:
logger.warning(
"Unconfirmed terminal finalization for run %s (%s); heartbeat RunStore.finalize_if_not_cancelled must support owner_worker_id and atomically enforce a live owner lease",
run_id,
type(exc).__name__,
)
await self._fence_unconfirmed_terminal(record)
return None
if result.finalized is True:
return None
if result.cancel_action is not None:
async with self._lock:
record = self._runs.get(run_id)
if record is not None:
record.abort_action = result.cancel_action
record.abort_event.set()
record.abort_action = result.cancel_action
record.abort_event.set()
return result.cancel_action
await self.set_status(
run_id,
status,
error=error,
stop_reason=stop_reason,
goal_verdict=goal_verdict,
persist=not result.finalized,
)
if result.finalized:
async with self._lock:
record = self._runs.get(run_id)
if record is not None:
record.terminal_status_persisted = True
# Rejected CAS is never retried through a weaker status/completion path.
await self._fence_unconfirmed_terminal(record)
return None
async def _ensure_delivery_receipt(self, record: RunRecord) -> bool:
@ -1431,13 +1650,14 @@ class RunManager:
"""Set process-local abort state without status persistence or cleanup."""
async with self._lock:
record = self._runs.get(run_id)
if record is None or (record.status not in (RunStatus.pending, RunStatus.running) and not record.scheduled_goal_cleanup_pending) or record.abort_event.is_set():
if record is None or (record.status not in (RunStatus.pending, RunStatus.running) and not self._awaits_terminal_commit(record)) or record.abort_event.is_set():
return
record.abort_action = action
record.abort_event.set()
task_active = record.task is not None and not record.task.done()
record.finalizing = task_active
record.rollback_finalization_pending = action == "rollback" and task_active
if task_active and record.status == RunStatus.running:
record.task.cancel()
logger.info("Run %s cancellation signalled locally (action=%s)", run_id, action)
@ -1511,6 +1731,7 @@ class RunManager:
record.abort_event.set()
task_active = record.task is not None and not record.task.done()
record.finalizing = task_active
record.rollback_finalization_pending = action == "rollback" and task_active
if task_active and record.status == RunStatus.running:
record.task.cancel()
record.status = RunStatus.interrupted
@ -1678,7 +1899,7 @@ class RunManager:
async def _close_cancelled_admission(self, record: RunRecord) -> None:
"""Terminalize an unseen replacement and confirm its durable state."""
await self.cancel(record.run_id)
if self._store is None:
if self._store is None or (self.heartbeat_enabled and record.ownership_lost):
return
stored = await self._call_store_with_retry(
@ -1692,11 +1913,16 @@ class RunManager:
# path needs a strict second CAS attempt because the caller never
# receives the record and no worker can attach after it returns.
# A peer terminal transition wins the CAS and is preserved below.
await self._call_store_with_retry(
"terminalize cancelled admission",
record.run_id,
lambda: self._store.update_status(record.run_id, RunStatus.interrupted.value),
)
if self.heartbeat_enabled:
await self._persist_status(record, RunStatus.interrupted)
if record.ownership_lost:
return
else:
await self._call_store_with_retry(
"terminalize cancelled admission",
record.run_id,
lambda: self._store.update_status(record.run_id, RunStatus.interrupted.value),
)
stored = await self._call_store_with_retry(
"verify terminal cancelled admission",
record.run_id,
@ -1909,10 +2135,14 @@ class RunManager:
max_retries = 3
for attempt in range(max_retries):
try:
await self._call_store_with_retry(
"create_thread_operation_atomic",
run_id,
lambda: self._store.create_thread_operation_atomic(**create_kwargs),
await self._atomic_replacement(
local_inflight,
multitask_strategy,
lambda: self._call_store_with_retry(
"create_thread_operation_atomic",
run_id,
lambda: self._store.create_thread_operation_atomic(**create_kwargs),
),
)
break
except RunIdempotencyConflict as exc:
@ -1935,10 +2165,11 @@ class RunManager:
# 3) Only now safe to register locally — store insert succeeded.
self._runs[run_id] = record
self._index_run_locked(record)
self._arm_lease_deadline(record)
# 4) Cancel local in-memory inflight (interrupt/rollback). The
# store-side counterparts were already cancelled in step 2.
if multitask_strategy in ("interrupt", "rollback"):
if self._store is None and multitask_strategy in ("interrupt", "rollback"):
for r in local_inflight:
if r.finalizing:
continue
@ -1957,6 +2188,9 @@ class RunManager:
# this point happens after the replacement was admitted, so close that
# new run before propagating cancellation to the caller.
try:
# Preserve the admission cancellation boundary without repeating
# a terminal write already acknowledged by the atomic transaction.
await asyncio.sleep(0)
for interrupted_record in interrupted_records:
await self._persist_status(interrupted_record, RunStatus.interrupted)
except asyncio.CancelledError:
@ -2270,27 +2504,36 @@ class RunManager:
expired, this worker is no longer authorized to publish a terminal
outcome. A peer reconciler owns durable terminalization.
"""
return self._fence_local_record(record, reason=reason, require_active=require_active)
def _fence_local_record(self, record: RunRecord, *, reason: str, require_active: bool = False) -> bool:
# Event-loop-local state transition, with no await. In particular this
# must not wait behind an admission holding _lock across Store I/O.
task_to_cancel: asyncio.Task | None = None
async with self._lock:
current = self._runs.get(record.run_id)
if current is not record:
current = self._runs.get(record.run_id)
if current is not record:
return False
if require_active:
if record.status not in (RunStatus.pending, RunStatus.running):
return False
if require_active:
if record.status not in (RunStatus.pending, RunStatus.running):
return False
if record.task is not None and record.task.done():
return False
if record.ownership_lost:
return True
record.ownership_lost = True
# A fenced worker must not keep the peer-recoverable row alive.
if record.task is not None and record.task.done():
return False
if record.ownership_lost:
return True
if self._has_final_terminal_ack(record):
record.terminal_commit_pending = False
record.abort_event.set()
record.status = RunStatus.error
record.error = reason
record.updated_at = _now_iso()
if record.task is not None and not record.task.done() and record.task is not asyncio.current_task():
task_to_cancel = record.task
return False
self._clear_lease_deadline(record.run_id)
record.ownership_lost = True
record.terminal_status_persisted = False
# A fenced worker must not keep the peer-recoverable row alive.
record.terminal_commit_pending = False
record.abort_event.set()
record.status = RunStatus.error
record.error = reason
record.updated_at = _now_iso()
if record.task is not None and not record.task.done() and record.task is not asyncio.current_task():
task_to_cancel = record.task
if task_to_cancel is not None:
task_to_cancel.cancel()
@ -2313,20 +2556,32 @@ class RunManager:
logger.info("Run lease heartbeat started for worker %s", self._worker_id)
async def stop_heartbeat(self, *, timeout: float = 5.0) -> None:
"""Stop the background heartbeat task within ``timeout`` seconds."""
"""Stop scheduling and boundedly drain owned lease checks and Store I/O."""
deadline = asyncio.get_running_loop().time() + max(0.0, timeout)
if self._heartbeat_stop is not None:
self._heartbeat_stop.set()
if self._heartbeat_task is not None and not self._heartbeat_task.done():
_, pending = await asyncio.wait(
(self._heartbeat_task,),
timeout=max(0.0, timeout),
)
if pending:
self._heartbeat_task.cancel()
try:
await self._heartbeat_task
except asyncio.CancelledError:
pass
def owned_tasks() -> set[asyncio.Task]:
tasks = set(self._lease_checks.values()) | set(self._lease_operations.values())
if self._heartbeat_task is not None:
tasks.add(self._heartbeat_task)
return {task for task in tasks if not task.done()}
tasks = owned_tasks()
while tasks and asyncio.get_running_loop().time() < deadline:
await asyncio.wait(tasks, timeout=max(0.0, deadline - asyncio.get_running_loop().time()))
# A check/heartbeat already in progress can register child I/O
# after the first snapshot. Include it in the same drain budget.
tasks = owned_tasks()
if tasks:
for task in tasks:
task.cancel()
# One scheduling turn delivers cooperative cancellations, without
# awaiting a Store that suppresses them beyond the drain deadline.
await asyncio.sleep(0)
unfinished = sum(not task.done() for task in tasks)
if unfinished:
logger.warning("Heartbeat stopped with %d lease task(s) still draining", unfinished)
self._heartbeat_task = None
self._heartbeat_stop = None
logger.info("Run lease heartbeat stopped for worker %s", self._worker_id)
@ -2359,7 +2614,7 @@ class RunManager:
cycle += 1
try:
await self._renew_leases()
await self._renew_leases(wait=False)
except Exception:
logger.warning("Heartbeat renewal cycle failed", exc_info=True)
@ -2375,188 +2630,218 @@ class RunManager:
@staticmethod
def _awaits_terminal_commit(record: RunRecord) -> bool:
"""Whether a run keeps its active durable row until a deferred terminal commit."""
return record.scheduled_goal_cleanup_pending or record.terminal_commit_pending
async def _terminal_write_confirmed(self, record: RunRecord) -> bool:
"""Recognize this worker's acknowledged outcome at the commit boundary.
A peer takeover may preserve owner_worker_id, so compare every terminal
outcome field rather than accepting any terminal row with that owner.
"""
if not self._awaits_terminal_commit(record) or record.status in (RunStatus.pending, RunStatus.running) or self._store is None:
return False
row = await self._store.get(record.run_id, user_id=record.user_id)
return (
row is not None
and row.get("thread_id") == record.thread_id
and row.get("user_id") == record.user_id
and row.get("owner_worker_id") == record.owner_worker_id
and row.get("status") == record.status.value
and row.get("error") == record.error
and row.get("stop_reason") == record.stop_reason
and row.get("goal_verdict") == record.goal_verdict
)
return record.scheduled_goal_cleanup_pending or record.terminal_commit_pending or record.rollback_finalization_pending
async def _release_confirmed_terminal_barrier(self, record: RunRecord, *, timeout: float) -> bool:
if not self._awaits_terminal_commit(record) or record.status in (RunStatus.pending, RunStatus.running):
return False
try:
async with asyncio.timeout(timeout):
if await self._terminal_write_confirmed(record):
record.scheduled_goal_cleanup_pending = False
record.terminal_commit_pending = False
return True
except Exception:
logger.warning("Unable to verify terminal commit for run %s", record.run_id, exc_info=True)
"""Wait only for local attempts, never infer attribution from a row read."""
deadline = asyncio.get_running_loop().time() + max(0.0, timeout)
while not record.ownership_lost:
if self._has_final_terminal_ack(record):
record.scheduled_goal_cleanup_pending = False
record.terminal_commit_pending = False
return True
attempts = set(self._terminal_attempts.get(record.run_id, ()))
remaining = deadline - asyncio.get_running_loop().time()
if not attempts or remaining <= 0:
return False
done, _ = await asyncio.wait(attempts, timeout=remaining, return_when=asyncio.FIRST_COMPLETED)
if not done:
# Check proof once more before treating the elapsed wait as
# unknown; an ACK can arrive on the timeout scheduling turn.
continue
return False
async def _renew_leases(self) -> None:
"""Renew locally-owned leases, failing closed at their deadlines.
def _clear_lease_deadline(self, run_id: str) -> None:
handle = self._lease_deadlines.pop(run_id, None)
if handle is not None:
handle.cancel()
``RunRecord.lease_expires_at`` advances only after a successful durable
renewal, so it is the last confirmed ownership deadline. Transient
exceptions are tolerated before that deadline; a call that blocks or
keeps failing through it fences the local run.
def _arm_lease_deadline(self, record: RunRecord) -> None:
self._clear_lease_deadline(record.run_id)
if not self.heartbeat_enabled or self._store is None or record.ownership_lost or self._has_final_terminal_ack(record):
return
expiry = self._parse_lease_deadline(record.lease_expires_at)
remaining = max(0.0, (expiry - datetime.now(UTC)).total_seconds()) if expiry else 0.0
def expired() -> None:
self._lease_deadlines.pop(record.run_id, None)
if self._runs.get(record.run_id) is not record or record.ownership_lost or self._has_final_terminal_ack(record):
return
current = self._parse_lease_deadline(record.lease_expires_at)
if current is not None and current > datetime.now(UTC):
self._arm_lease_deadline(record)
return
self._fence_local_record(record, reason="The last confirmed lease expired without a final terminal acknowledgment.")
# This local-only guard survives stopped/failed heartbeat I/O. Shutdown
# is not authority to continue execution past the last confirmed lease.
self._lease_deadlines[record.run_id] = asyncio.get_running_loop().call_later(remaining, expired)
def _supervise_lease_task(self, registry: dict, run_id: str, task: asyncio.Task) -> None:
registry[run_id] = task
def settled(completed: asyncio.Task) -> None:
if registry.get(run_id) is completed:
registry.pop(run_id, None)
if not completed.cancelled():
exc = completed.exception()
if exc is not None:
logger.warning("Lease task for run %s failed (%s)", run_id, type(exc).__name__)
task.add_done_callback(settled)
async def _renew_leases(self, *, wait: bool = True) -> None:
"""Give every run an independent renewal decision and monotonic budget.
The heartbeat starts checks without joining them, so its next tick can
renew healthy runs while a different run is still waiting for Store I/O.
Explicit callers may join the bounded decisions for one snapshot.
"""
if self._store is None or self._run_ownership_config is None:
return
lease_seconds = self._run_ownership_config.lease_seconds
cancellations: list[tuple[str, str]] = []
async with self._lock:
# Renew any pending/running run owned by this worker unless its
# background task has already completed. A pending run whose task
# has not been spawned yet (``task is None``) is still live from
# this worker's perspective — between ``create_thread_operation_atomic``
# inserting the row and the worker layer spawning the agent task
# there is a brief window. If we drop those records here and the
# window stretches past ``lease_seconds`` (e.g. event-loop
# saturation, slow checkpoint hydrate on a fresh worker), peer
# reconciliation will reclaim the run as an orphan and mark it
# ``error`` even though this worker still intends to execute it.
# The same holds for a locally terminal run whose worker is still
# finalizing before its deferred terminal commit.
active_runs = [
(rid, record)
for rid, record in self._runs.items()
if (record.status in (RunStatus.pending, RunStatus.running) or self._awaits_terminal_commit(record)) and record.owner_worker_id == self._worker_id and (record.task is None or not record.task.done())
records = [
record
for record in self._runs.values()
if not record.ownership_lost
and record.owner_worker_id == self._worker_id
and (record.status in (RunStatus.pending, RunStatus.running) or self._awaits_terminal_commit(record))
and (self._awaits_terminal_commit(record) or record.task is None or not record.task.done())
]
checks = []
for record in records:
task = self._lease_checks.get(record.run_id)
if task is None or task.done():
task = asyncio.create_task(self._renew_one_lease(record), name=f"run-lease-check:{record.run_id}")
self._supervise_lease_task(self._lease_checks, record.run_id, task)
checks.append(task)
if wait and checks:
# Cancellation of a cycle must not detach or cancel another caller's
# still-owned per-run decision. Shutdown drains the registries.
await asyncio.wait(checks)
for run_id, record in active_runs:
if record.status not in (RunStatus.pending, RunStatus.running) and not self._awaits_terminal_commit(record):
continue
confirmed_deadline = self._parse_lease_deadline(record.lease_expires_at)
if self._awaits_terminal_commit(record) and record.status not in (RunStatus.pending, RunStatus.running):
# The DB may have committed the terminal CAS while the worker
# is still receiving its acknowledgement. It no longer needs
# renewal and must not be mistaken for a peer takeover.
remaining = (confirmed_deadline - datetime.now(UTC)).total_seconds() if confirmed_deadline is not None else 0
if await self._release_confirmed_terminal_barrier(record, timeout=remaining if remaining > 0 else lease_seconds):
continue
if confirmed_deadline is None or confirmed_deadline <= datetime.now(UTC):
if await self._release_confirmed_terminal_barrier(record, timeout=lease_seconds):
continue
await self._mark_ownership_lost(
record,
reason="Lease ownership could not be confirmed before the last confirmed lease expired.",
require_active=not self._awaits_terminal_commit(record),
)
continue
def _rollback_retry_authorized(self, record: RunRecord, *, initial_allow_interrupted: bool) -> bool:
"""Whether an active-only rejection may be retried for this rollback."""
proof = record.terminal_commit_proof
return (
not initial_allow_interrupted
and record.rollback_finalization_pending
and proof is not None
and proof.status == RunStatus.interrupted
and record.owner_worker_id == self._worker_id
and self._runs.get(record.run_id) is record
and not record.ownership_lost
and not self._has_final_terminal_ack(record)
)
remaining = (confirmed_deadline - datetime.now(UTC)).total_seconds()
new_expiry = (datetime.now(UTC) + timedelta(seconds=lease_seconds)).isoformat()
try:
async with asyncio.timeout(remaining):
renewal = await self._call_store_with_retry(
"renew_lease",
run_id,
lambda: self._store.renew_lease(
run_id,
owner_worker_id=self._worker_id,
lease_expires_at=new_expiry,
),
)
if renewal.renewed:
if confirmed_deadline <= datetime.now(UTC):
if await self._release_confirmed_terminal_barrier(record, timeout=lease_seconds):
continue
await self._mark_ownership_lost(
record,
reason="Lease renewal completed after the last confirmed lease had already expired.",
require_active=not self._awaits_terminal_commit(record),
)
continue
# Unsynced write is benign: ``lease_expires_at`` is the
# only field on an existing record this path mutates, so
# there is no concurrent writer to race against
# (``set_status`` / ``_persist_status`` touch other
# fields). Re-acquiring ``self._lock`` here would
# serialise against unrelated run mutations for no gain.
record.lease_expires_at = new_expiry
if renewal.cancel_action is not None:
action = renewal.cancel_action
if action not in ("interrupt", "rollback"):
logger.warning(
"Run %s has invalid durable cancel action %r; using interrupt",
run_id,
action,
)
action = "interrupt"
cancellations.append((run_id, action))
else:
# ``renew_lease`` returned False — the row was claimed
# by another worker (status is no longer pending/running,
# or ``owner_worker_id`` changed). Stop the local task so
# we don't waste CPU or overwrite the takeover status on
# finalisation.
async with self._lock:
still_active = (
self._runs.get(run_id) is record
and (record.status in (RunStatus.pending, RunStatus.running) or self._awaits_terminal_commit(record))
and record.owner_worker_id == self._worker_id
and (record.task is None or not record.task.done())
)
if still_active:
# A detached precheck may have observed running just
# before our terminal CAS committed. Re-read before
# treating the rejected renewal as lost ownership.
if await self._release_confirmed_terminal_barrier(record, timeout=lease_seconds):
continue
logger.warning(
"Run %s lease renewal failed (status=%s,owner=%s) – worker likely taken over; aborting local task",
run_id,
record.status.value,
record.owner_worker_id,
)
await self._mark_ownership_lost(
record,
reason="The durable store rejected lease renewal for this worker.",
require_active=not self._awaits_terminal_commit(record),
)
except Exception:
if confirmed_deadline <= datetime.now(UTC):
if await self._release_confirmed_terminal_barrier(record, timeout=lease_seconds):
continue
await self._mark_ownership_lost(
record,
reason="Lease ownership could not be confirmed before the last confirmed lease expired.",
require_active=not self._awaits_terminal_commit(record),
)
else:
logger.warning(
"Failed to renew lease for run %s before its confirmed deadline; will retry",
run_id,
exc_info=True,
)
async def _await_rollback_retry_authority(self, record: RunRecord, *, initial_allow_interrupted: bool, deadline: float) -> bool:
"""Wait for an in-flight admission ACK, never infer it from the row."""
loop = asyncio.get_running_loop()
while loop.time() < deadline:
if self._rollback_retry_authorized(record, initial_allow_interrupted=initial_allow_interrupted):
return True
attempts = set(self._terminal_attempts.get(record.run_id, ()))
remaining = deadline - loop.time()
if not attempts or remaining <= 0:
break
await asyncio.wait(attempts, timeout=remaining, return_when=asyncio.FIRST_COMPLETED)
return loop.time() < deadline and self._rollback_retry_authorized(record, initial_allow_interrupted=initial_allow_interrupted)
# Keep cancellation status writes and cleanup out of the sole renewal
# loop. After every local lease has had a chance to renew, only signal
# the owning worker task; that task performs normal terminal handling.
for run_id, action in cancellations:
await self._signal_local_cancel(
run_id,
action=action,
async def _renew_one_lease(self, record: RunRecord) -> None:
assert self._store is not None and self._run_ownership_config is not None
if self._has_final_terminal_ack(record):
await self._release_confirmed_terminal_barrier(record, timeout=0)
return
self._arm_lease_deadline(record)
confirmed = self._parse_lease_deadline(record.lease_expires_at)
remaining = max(0.0, (confirmed - datetime.now(UTC)).total_seconds()) if confirmed else 0.0
loop = asyncio.get_running_loop()
deadline = loop.time() + remaining
async def fence(reason: str) -> None:
if not await self._release_confirmed_terminal_barrier(record, timeout=max(0.0, deadline - loop.time())):
await self._mark_ownership_lost(record, reason=reason, require_active=False)
if remaining <= 0:
await fence("Lease ownership could not be confirmed before the last confirmed lease expired.")
return
if self._awaits_terminal_commit(record) and record.task is not None and record.task.done() and not self._terminal_attempts.get(record.run_id):
await fence("The worker exited without an attributable terminal acknowledgment.")
return
if record.run_id in self._lease_operations:
# A cancellation-resistant operation is still owned and observed.
# Never start another renewal for this run while it is in flight.
await fence("A previous lease renewal is still unconfirmed.")
return
initial_allow_interrupted = record.rollback_finalization_pending and record.terminal_commit_proof is not None and record.terminal_commit_proof.status == RunStatus.interrupted
for retry_index in range(2):
allow_interrupted = record.rollback_finalization_pending and record.terminal_commit_proof is not None and record.terminal_commit_proof.status == RunStatus.interrupted
new_expiry = (datetime.now(UTC) + timedelta(seconds=self._run_ownership_config.lease_seconds)).isoformat()
operation = asyncio.create_task(
self._call_store_with_retry(
"renew_lease",
record.run_id,
lambda: self._store.renew_lease(
record.run_id,
owner_worker_id=self._worker_id,
lease_expires_at=new_expiry,
**({"allow_interrupted": True} if allow_interrupted else {}),
),
),
name=f"run-lease-io:{record.run_id}",
)
self._supervise_lease_task(self._lease_operations, record.run_id, operation)
try:
done, _ = await asyncio.wait({operation}, timeout=max(0.0, deadline - loop.time()))
if not done:
operation.cancel()
# Do not join cancellation-resistant I/O. Its late result is
# observed by the supervisor and cannot mutate local authority.
await fence("Lease renewal exceeded the last confirmed lease deadline.")
return
try:
renewal = operation.result()
except (Exception, asyncio.CancelledError):
if loop.time() >= deadline:
await fence("Lease renewal failed through the last confirmed lease deadline.")
else:
logger.warning("Lease renewal failed for run %s; will retry before expiry", record.run_id)
return
if self._has_final_terminal_ack(record):
await self._release_confirmed_terminal_barrier(record, timeout=0)
return
if record.ownership_lost or self._runs.get(record.run_id) is not record:
return
if loop.time() >= deadline:
await fence("Lease renewal completed after the last confirmed lease had already expired.")
return
if renewal.renewed is True:
record.lease_expires_at = new_expiry
self._arm_lease_deadline(record)
if renewal.cancel_action is not None:
action = renewal.cancel_action if renewal.cancel_action in ("interrupt", "rollback") else "interrupt"
# Signal here, before joining any other run's decision.
await self._signal_local_cancel(record.run_id, action=action)
return
if (
retry_index == 0
and not initial_allow_interrupted
and await self._await_rollback_retry_authority(
record,
initial_allow_interrupted=initial_allow_interrupted,
deadline=deadline,
)
):
# The first request raced atomic rollback admission. Re-check
# the interrupted row once under the same per-run decision,
# with owner and live-lease predicates still enforced in Store.
continue
await fence("The durable store rejected lease renewal for this worker.")
return
except asyncio.CancelledError:
operation.cancel()
raise
async def _reconcile_orphans_periodic(self) -> None:
"""Sweep for expired leases owned by dead peers.
@ -2617,9 +2902,13 @@ class RunManager:
_, pending = await asyncio.wait((task,), timeout=max(0.0, timeout))
if pending:
task.cancel()
await asyncio.gather(task, return_exceptions=True)
await asyncio.sleep(0) # Deliver cancellation without joining I/O.
if task.done():
self._orphan_reconciliation_done(task)
# Retain the single-flight reference and done callback. Cancellation
# is a request; a Store may keep using its resources after return.
logger.warning(
"Orphan recovery drain exceeded %.1fs on shutdown; cancelled the active pass",
"Orphan recovery drain exceeded %.1fs on shutdown; cancellation requested, Store I/O may still be active",
timeout,
)

43
backend/packages/harness/deerflow/runtime/runs/store/base.py
View File

@ -127,6 +127,12 @@ class RunStore(abc.ABC):
# Inheriting the opted-in method unchanged preserves its capability.
supports_idempotency_request_kwargs = False
# Strict completion requires explicit require_terminal and owner_worker_id
# keywords, or this opt-in on the class defining update_run_completion.
# Inheriting that method preserves support; overriding it must opt in again.
# Accepting **kwargs alone never promises the atomic terminal/owner guards.
supports_terminal_completion_kwargs = False
async def list_changed(
self,
*,
@ -235,11 +241,13 @@ class RunStore(abc.ABC):
error: str | None = None,
stop_reason: str | None = None,
goal_verdict: dict[str, Any] | None = None,
owner_worker_id: str | None = None,
) -> bool | None:
"""Update a run status.
Returns ``False`` when the store can prove no row was updated. Older or
lightweight stores may return ``None`` when they cannot report rowcount.
With ``owner_worker_id``, atomically require that owner and an unexpired
lease. Return True only after this operation commits. False means no
update; None from legacy stores is unconfirmed in heartbeat mode.
"""
pass
@ -279,6 +287,8 @@ class RunStore(abc.ABC):
run_id: str,
*,
status: str,
require_terminal: bool = False,
owner_worker_id: str | None = None,
total_input_tokens: int = 0,
total_output_tokens: int = 0,
total_tokens: int = 0,
@ -297,7 +307,9 @@ class RunStore(abc.ABC):
Implementations must not replace a different terminal status. Returns
``False`` when the row is missing or already has a conflicting terminal
outcome.
outcome. With ``require_terminal``, require the exact requested terminal
status and owner atomically: metadata must not transition active rows or
change error/goal_verdict. Unsupported strict calls must fail closed.
"""
pass
@ -346,8 +358,14 @@ class RunStore(abc.ABC):
*,
owner_worker_id: str,
lease_expires_at: str,
allow_interrupted: bool = False,
) -> bool:
"""Renew the lease on an active run. Returns ``False`` when no row matched."""
"""Renew only an active, owned, unexpired lease to a future deadline.
The explicit allow_interrupted option also permits an interrupted row
while its owning Worker finishes rollback. Never renew success/error.
Return False when no row satisfies these atomic conditions.
"""
pass
async def renew_lease(
@ -356,6 +374,7 @@ class RunStore(abc.ABC):
*,
owner_worker_id: str,
lease_expires_at: str,
allow_interrupted: bool = False,
) -> LeaseRenewal:
"""Renew ownership and return any durable cancellation request.
@ -363,12 +382,15 @@ class RunStore(abc.ABC):
cancellation action, so third-party stores remain source-compatible
without adding a background read. Stores that support multi-process
cancellation must override this method to renew and observe the
request atomically.
request atomically, retaining the live-lease/owner conditions. The
allow_interrupted option must retain those predicates; unsupported
implementations fail closed rather than silently relaxing status guards.
"""
renewed = await self.update_lease(
run_id,
owner_worker_id=owner_worker_id,
lease_expires_at=lease_expires_at,
**({"allow_interrupted": True} if allow_interrupted else {}),
)
return LeaseRenewal(renewed=renewed)
@ -388,12 +410,17 @@ class RunStore(abc.ABC):
error: str | None = None,
stop_reason: str | None = None,
goal_verdict: dict[str, Any] | None = None,
owner_worker_id: str | None = None,
) -> StatusFinalization:
"""Atomically finalize an active run unless cancellation won.
The compatibility default is safe for stores that do not implement
durable cancellation.
With an owner, atomically require a live owner lease and return an
explicit commit acknowledgment. Rejection without cancel_action is final,
not permission to retry a weaker write. Exceptions/cancellation are unknown.
The legacy default cannot enforce this and rejects owned calls.
"""
if owner_worker_id is not None:
raise NotImplementedError("Store must implement owner/lease-fenced finalization")
updated = await self.update_status(
run_id,
status,
@ -401,7 +428,7 @@ class RunStore(abc.ABC):
stop_reason=stop_reason,
**({"goal_verdict": goal_verdict} if goal_verdict is not None else {}),
)
return StatusFinalization(finalized=updated is not False)
return StatusFinalization(finalized=updated is True)
@abc.abstractmethod
async def claim_for_takeover(

36
backend/packages/harness/deerflow/runtime/runs/store/memory.py
View File

@ -205,10 +205,16 @@ class MemoryRunStore(RunStore):
thread_run_ids = self._runs_by_thread.get(thread_id) or ()
return {run_id: run for run_id in thread_run_ids if run_id in run_ids and (run := self._runs.get(run_id)) is not None and run.get("operation_kind", "run") == "run" and (user_id is None or run.get("user_id") == user_id)}
async def update_status(self, run_id, status, *, error=None, stop_reason=None, goal_verdict=None):
async def update_status(self, run_id, status, *, error=None, stop_reason=None, goal_verdict=None, owner_worker_id=None):
run = self._runs.get(run_id)
if run is None:
return False
if owner_worker_id is not None and not self._live_owner(run, owner_worker_id):
return False
if owner_worker_id is not None and run["status"] == "interrupted" and status != "error":
return False
if owner_worker_id is not None and status == "success" and (run["status"] not in ("pending", "running") or run.get("cancel_action") is not None):
return False
# Guard: only transition rows that are still active. ``interrupted``
# is included for the rollback path (``interrupted → error`` finalize).
if run["status"] not in ("pending", "running", "interrupted"):
@ -263,11 +269,13 @@ class MemoryRunStore(RunStore):
removed += 1
return removed
async def update_run_completion(self, run_id, *, status, **kwargs):
async def update_run_completion(self, run_id, *, status, require_terminal=False, owner_worker_id=None, **kwargs):
run = self._runs.get(run_id)
if run is None:
return False
current_status = run.get("status")
if require_terminal and (status in ("pending", "running") or current_status != status or run.get("owner_worker_id") != owner_worker_id):
return False
allowed_sources = {"pending", "running", status}
if status == "error":
allowed_sources.add("interrupted")
@ -275,6 +283,8 @@ class MemoryRunStore(RunStore):
return False
run["status"] = status
for key, value in kwargs.items():
if require_terminal and key in ("error", "stop_reason", "goal_verdict"):
continue
if value is not None:
run[key] = copy.deepcopy(value) if key == "goal_verdict" else value
run["updated_at"] = datetime.now(UTC).isoformat()
@ -340,19 +350,28 @@ class MemoryRunStore(RunStore):
# Multi-worker run ownership methods
# ------------------------------------------------------------------
@staticmethod
def _live_owner(run, owner_worker_id) -> bool:
from deerflow.utils.time import is_lease_expired
return run.get("owner_worker_id") == owner_worker_id and not is_lease_expired(run.get("lease_expires_at"), grace_seconds=0)
async def update_lease(
self,
run_id: str,
*,
owner_worker_id: str,
lease_expires_at: str,
allow_interrupted: bool = False,
) -> bool:
run = self._runs.get(run_id)
if run is None:
return False
if run["status"] not in ("pending", "running"):
if run["status"] not in (("pending", "running", "interrupted") if allow_interrupted else ("pending", "running")):
return False
if run.get("owner_worker_id") != owner_worker_id:
if not self._live_owner(run, owner_worker_id):
return False
if datetime.fromisoformat(lease_expires_at) <= datetime.now(UTC):
return False
run["owner_worker_id"] = owner_worker_id
run["lease_expires_at"] = lease_expires_at
@ -365,6 +384,7 @@ class MemoryRunStore(RunStore):
*,
owner_worker_id: str,
lease_expires_at: str,
allow_interrupted: bool = False,
) -> LeaseRenewal:
# Delegate through ``update_lease`` so lightweight subclasses and tests
# that override the legacy primitive keep the same behavior.
@ -372,6 +392,7 @@ class MemoryRunStore(RunStore):
run_id,
owner_worker_id=owner_worker_id,
lease_expires_at=lease_expires_at,
**({"allow_interrupted": True} if allow_interrupted else {}),
)
if not renewed:
return LeaseRenewal(renewed=False)
@ -402,10 +423,13 @@ class MemoryRunStore(RunStore):
error: str | None = None,
stop_reason: str | None = None,
goal_verdict: dict[str, Any] | None = None,
owner_worker_id: str | None = None,
) -> StatusFinalization:
run = self._runs.get(run_id)
if run is None:
return StatusFinalization(finalized=False)
if owner_worker_id is not None and (run["status"] not in ("pending", "running") or not self._live_owner(run, owner_worker_id)):
return StatusFinalization(finalized=False)
if run.get("cancel_action") is not None:
return StatusFinalization(
finalized=False,
@ -613,4 +637,6 @@ class MemoryRunStore(RunStore):
new_row["change_seq"] = change_seq
self._runs[run_id] = new_row
self._index_run(run_id, thread_id)
return new_row, claimed
# Commit acknowledgments must not change when a later writer mutates
# the in-memory row while the caller is waiting to receive this result.
return copy.deepcopy(new_row), copy.deepcopy(claimed)

27
backend/packages/harness/deerflow/runtime/runs/worker.py
View File

@ -1019,14 +1019,21 @@ async def run_agent(
restore_checkpoint: bool = True,
) -> None:
nonlocal checkpoint_rollback_completed
if record.ownership_lost:
return
await run_manager.set_finalizing(run_id, True)
if record.ownership_lost:
return
if action == "rollback":
record.rollback_finalization_pending = run_manager.heartbeat_enabled and run_manager._store is not None
await run_manager.set_status(
run_id,
RunStatus.error,
error="Rolled back by user",
**terminal_status_kwargs,
**({"persist": False} if record.rollback_finalization_pending else terminal_status_kwargs),
)
if record.ownership_lost:
return
if not restore_checkpoint:
return
try:
@ -1811,17 +1818,17 @@ async def run_agent(
logger.warning("Could not remove scheduled goal for run %s", run_id, exc_info=True)
await run_manager.set_status(run_id, RunStatus.error, error="Scheduled goal cleanup failed", persist=False)
if not record.ownership_lost and (event_store is not None or scheduled_goal is not None):
if not record.ownership_lost and (event_store is not None or scheduled_goal is not None or record.rollback_finalization_pending):
try:
# Even after bounded receipt retries are exhausted, persist the
# real worker outcome. Leaving a successful row inflight would
# let lease recovery rewrite it as an error with a synthetic
# zero receipt.
if record.abort_event.is_set():
if scheduled_goal is not None:
# A cancel may arrive while scheduled cleanup was
# awaiting its saver. The local status may still be
# staged success; apply the accepted action, including
if scheduled_goal is not None or record.status == RunStatus.success:
# A cancel may arrive while receipt or scheduled
# cleanup is awaiting I/O. Apply the accepted action
# to a still-staged success, including
# rollback, before committing any terminal outcome.
await _finish_cancellation(record.abort_action, restore_checkpoint=not checkpoint_rollback_completed)
await run_manager.persist_current_status(run_id)
@ -1832,9 +1839,11 @@ async def run_agent(
await run_manager.persist_current_status(run_id)
except Exception:
logger.warning("Failed to persist terminal status for run %s after delivery receipt attempts", run_id, exc_info=True)
# The deferred commit has been attempted. A failed write is left to
# completion fallback or lease recovery, as before.
record.terminal_commit_pending = False
# Only a confirmed terminal write or ownership loss releases the
# renewal barrier in multi-worker mode. A returned/aborted call is
# not evidence that this worker committed its terminal outcome.
if record.terminal_commit_pending and run_manager.can_release_terminal_barrier(record):
record.terminal_commit_pending = False
if not record.ownership_lost and journal is not None and persist_completion:
try:

387
backend/tests/test_atomic_admission_authority.py Normal file
View File

@ -0,0 +1,387 @@
import asyncio
from datetime import UTC, datetime, timedelta
from unittest.mock import AsyncMock
import pytest
from test_terminal_store_contract import store as store
from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.runtime import RunManager, RunStatus
@pytest.fixture
def anyio_backend():
return "asyncio"
def manager_for(store):
return RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=5, grace_seconds=0))
@pytest.mark.anyio
@pytest.mark.parametrize("strategy", ["interrupt", "rollback"])
async def test_same_owner_atomic_replacement_does_not_fence_its_old_run(store, strategy):
manager = manager_for(store)
first = await manager.create_or_reject("thread", user_id="review")
await manager.try_start(first.run_id)
second = await manager.create_or_reject("thread", multitask_strategy=strategy, user_id="review")
row = await store.get(first.run_id, user_id="review")
assert row["owner_worker_id"] == "owner"
assert row["status"] == "interrupted"
assert datetime.fromisoformat(row["lease_expires_at"]) > datetime.now(UTC)
assert second.status == RunStatus.pending
if strategy == "rollback":
await manager.set_status(first.run_id, RunStatus.error, error="Rolled back by user")
row = await store.get(first.run_id, user_id="review")
assert row["status"] == "error"
assert row["error"] == "Rolled back by user"
assert not first.ownership_lost, "successful same-owner atomic replacement was treated as lost ownership"
@pytest.mark.anyio
async def test_fast_renewal_exception_retains_deadline_fencing():
from deerflow.runtime.runs.store.memory import MemoryRunStore
store = MemoryRunStore()
manager = manager_for(store)
record = await manager.create_or_reject("thread")
await manager.try_start(record.run_id)
record.task = asyncio.create_task(asyncio.Event().wait())
expiry = datetime.now(UTC) + timedelta(milliseconds=100)
record.lease_expires_at = expiry.isoformat()
store._runs[record.run_id]["lease_expires_at"] = expiry.isoformat()
store.renew_lease = AsyncMock(side_effect=RuntimeError("unavailable"))
expired = asyncio.Event()
loop = asyncio.get_running_loop()
handle = loop.call_later(0.15, expired.set)
await manager.start_heartbeat()
try:
await manager._renew_leases()
await expired.wait()
assert datetime.now(UTC) > expiry
assert await store.claim_for_takeover(record.run_id, grace_seconds=0, error="peer recovery")
assert record.ownership_lost, "the failed lease check exited and left no deadline decision before the next heartbeat tick"
finally:
handle.cancel()
await manager.stop_heartbeat(timeout=1)
record.task.cancel()
await asyncio.gather(record.task, return_exceptions=True)
@pytest.mark.anyio
async def test_shutdown_does_not_join_uncooperative_orphan_io_past_budget():
from deerflow.runtime.runs.store.memory import MemoryRunStore
store = MemoryRunStore()
manager = manager_for(store)
entered, release, cancelled = asyncio.Event(), asyncio.Event(), asyncio.Event()
async def resistant(*args, **kwargs):
entered.set()
while not release.is_set():
try:
await release.wait()
except asyncio.CancelledError:
cancelled.set()
return []
store.list_inflight_with_expired_lease = resistant
manager._schedule_orphan_reconciliation()
recovery = manager._orphan_recovery_task
await asyncio.wait_for(entered.wait(), 1)
shutdown = asyncio.create_task(manager.shutdown(timeout=0.01))
try:
await asyncio.wait_for(cancelled.wait(), 1)
done, _ = await asyncio.wait({shutdown}, timeout=0.05)
assert shutdown in done, "shutdown joined cancellation-resistant orphan I/O without a bound"
assert manager._orphan_recovery_task is recovery
assert not recovery.done()
manager._schedule_orphan_reconciliation()
assert manager._orphan_recovery_task is recovery
finally:
release.set()
await asyncio.wait_for(shutdown, 1)
await asyncio.wait_for(recovery, 1)
@pytest.mark.anyio
async def test_cancelled_admission_compensation_cannot_write_after_expiry(store):
manager = manager_for(store)
record = await manager.create_or_reject("thread", user_id="review")
expired = (datetime.now(UTC) - timedelta(seconds=1)).isoformat()
await store.put(record.run_id, thread_id="thread", status="pending", user_id="review", owner_worker_id="owner", lease_expires_at=expired)
record.lease_expires_at = expired
await manager._close_cancelled_admission(record)
row = await store.get(record.run_id, user_id="review")
assert row["status"] == "pending", "unowned compensation terminalized a row after the owned write had lost its lease"
@pytest.mark.anyio
async def test_atomic_attempt_is_registered_until_its_actual_claim_ack():
from deerflow.runtime.runs.store.memory import MemoryRunStore
store = MemoryRunStore()
manager = manager_for(store)
old = await manager.create_or_reject("thread")
await manager.try_start(old.run_id)
committed, release = asyncio.Event(), asyncio.Event()
original = store.create_thread_operation_atomic
async def hold_ack(**kwargs):
assert manager._terminal_attempts.get(old.run_id), "attempt must precede Store I/O"
result = await original(**kwargs)
committed.set()
await release.wait()
return result
store.create_thread_operation_atomic = hold_ack
admission = asyncio.create_task(manager.create_or_reject("thread", multitask_strategy="interrupt"))
try:
await asyncio.wait_for(committed.wait(), 1)
assert manager._terminal_attempts.get(old.run_id)
assert not manager._has_terminal_ack(old)
release.set()
await asyncio.wait_for(admission, 1)
assert manager._has_terminal_ack(old)
assert old.error == "Cancelled by newer run"
assert not old.ownership_lost
assert not manager._terminal_attempts
finally:
release.set()
await asyncio.gather(admission, return_exceptions=True)
await manager.stop_heartbeat(timeout=0.1)
@pytest.mark.anyio
async def test_unclaimed_local_row_cannot_borrow_atomic_admission_ack():
from deerflow.runtime.runs.store.memory import MemoryRunStore
store = MemoryRunStore()
manager = manager_for(store)
old = await manager.create_or_reject("thread")
await manager.try_start(old.run_id)
original = store.create_thread_operation_atomic
async def omit_claim(**kwargs):
new, _ = await original(**kwargs)
return new, []
store.create_thread_operation_atomic = omit_claim
await manager.create_or_reject("thread", multitask_strategy="rollback")
assert old.status != RunStatus.interrupted
assert old.terminal_commit_proof is None
assert not manager._terminal_attempts
await manager.stop_heartbeat(timeout=0.1)
@pytest.mark.anyio
async def test_atomic_exception_settles_attempt_without_cancelling_old_worker():
from deerflow.runtime.runs.store.memory import MemoryRunStore
store = MemoryRunStore()
manager = manager_for(store)
old = await manager.create_or_reject("thread")
async def reject(**kwargs):
assert manager._terminal_attempts.get(old.run_id)
raise RuntimeError("admission rejected")
store.create_thread_operation_atomic = reject
with pytest.raises(RuntimeError, match="admission rejected"):
await manager.create_or_reject("thread", multitask_strategy="rollback")
assert not manager._terminal_attempts
assert old.terminal_commit_proof is None
assert not old.abort_event.is_set()
assert old.status == RunStatus.pending
await manager.stop_heartbeat(timeout=0.1)
@pytest.mark.anyio
async def test_successful_renewal_invalidates_the_earlier_deadline_guard():
from deerflow.runtime.runs.store.memory import MemoryRunStore
store = MemoryRunStore()
manager = manager_for(store)
record = await manager.create_or_reject("thread")
initial = (datetime.now(UTC) + timedelta(milliseconds=100)).isoformat()
record.lease_expires_at = initial
store._runs[record.run_id]["lease_expires_at"] = initial
original = store.renew_lease
store.renew_lease = AsyncMock(side_effect=RuntimeError("unavailable"))
await manager._renew_leases()
store.renew_lease = original
await manager._renew_leases()
elapsed = asyncio.Event()
asyncio.get_running_loop().call_later(0.15, elapsed.set)
await elapsed.wait()
assert not record.ownership_lost
assert record.lease_expires_at != initial
await manager.stop_heartbeat(timeout=0.1)
@pytest.mark.anyio
async def test_expiry_fences_even_while_atomic_ack_holds_the_manager_lock():
from deerflow.runtime.runs.store.memory import MemoryRunStore
store = MemoryRunStore()
manager = manager_for(store)
old = await manager.create_or_reject("thread")
original_atomic, original_renew = store.create_thread_operation_atomic, store.renew_lease
expiry = (datetime.now(UTC) + timedelta(milliseconds=100)).isoformat()
old.lease_expires_at = expiry
store._runs[old.run_id]["lease_expires_at"] = expiry
renewing, allow_renew, committed, release_ack = (asyncio.Event() for _ in range(4))
async def held_renew(*args, **kwargs):
renewing.set()
await allow_renew.wait()
return await original_renew(*args, **kwargs)
async def held_atomic(**kwargs):
result = await original_atomic(**kwargs)
committed.set()
await release_ack.wait()
return result
store.renew_lease, store.create_thread_operation_atomic = held_renew, held_atomic
renewal = asyncio.create_task(manager._renew_leases())
admission = None
try:
await asyncio.wait_for(renewing.wait(), 1)
admission = asyncio.create_task(manager.create_or_reject("thread", multitask_strategy="rollback"))
await asyncio.wait_for(committed.wait(), 1)
allow_renew.set()
done, _ = await asyncio.wait({renewal}, timeout=0.3)
assert renewal in done, "expiry fencing waited behind the atomic Store await under the manager lock"
assert old.ownership_lost
release_ack.set()
await asyncio.wait_for(admission, 1)
assert old.ownership_lost
assert not manager._has_terminal_ack(old)
assert not manager._terminal_attempts
finally:
allow_renew.set()
release_ack.set()
await asyncio.gather(renewal, *([admission] if admission else []), return_exceptions=True)
await manager.stop_heartbeat(timeout=0.1)
@pytest.mark.anyio
@pytest.mark.parametrize("method", ["renew_lease", "update_lease"])
@pytest.mark.parametrize(
"status,owner,expired,allowed",
[
("interrupted", "owner", False, True),
("interrupted", "peer", False, False),
("interrupted", "owner", True, False),
("error", "owner", False, False),
("success", "owner", False, False),
],
)
async def test_interrupted_renewal_requires_opt_in_and_live_owner(store, method, status, owner, expired, allowed):
expiry = (datetime.now(UTC) + timedelta(seconds=-1 if expired else 30)).isoformat()
await store.put("run", thread_id="thread", status=status, user_id=None, owner_worker_id="owner", lease_expires_at=expiry)
kwargs = dict(owner_worker_id=owner, lease_expires_at=(datetime.now(UTC) + timedelta(seconds=60)).isoformat())
result = await getattr(store, method)("run", **kwargs)
assert not (result.renewed if method == "renew_lease" else result)
result = await getattr(store, method)("run", allow_interrupted=True, **kwargs)
assert (result.renewed if method == "renew_lease" else result) is allowed
@pytest.mark.anyio
async def test_atomic_claim_ack_is_a_snapshot_of_the_committed_outcome(store):
manager = manager_for(store)
old = await manager.create_or_reject("thread", user_id="review")
_, claims = await store.create_thread_operation_atomic(
run_id="replacement",
thread_id="thread",
owner_worker_id="owner",
lease_expires_at=old.lease_expires_at,
multitask_strategy="rollback",
user_id="review",
)
assert len(claims) == 1
assert claims[0]["run_id"] == old.run_id
assert claims[0]["owner_worker_id"] == "owner"
assert claims[0]["status"] == "interrupted"
assert claims[0]["error"] == "Cancelled by newer run"
assert await store.update_status(old.run_id, "error", error="Rolled back by user", owner_worker_id="owner")
assert claims[0]["status"] == "interrupted"
assert claims[0]["error"] == "Cancelled by newer run"
@pytest.mark.anyio
async def test_atomic_replacement_rejects_live_peer_without_ack(store):
from deerflow.runtime.runs.manager import ConflictError
manager = manager_for(store)
old = await manager.create_or_reject("thread", user_id="review")
peer = RunManager(store=store, worker_id="peer", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=5, grace_seconds=0))
with pytest.raises(ConflictError):
await peer.create_or_reject("thread", multitask_strategy="rollback", user_id="review")
assert (await store.get(old.run_id, user_id="review"))["status"] == "pending"
assert old.terminal_commit_proof is None
assert not old.abort_event.is_set()
assert not peer._terminal_attempts
@pytest.mark.anyio
async def test_pending_caller_cancellation_closes_admitted_replacement():
from deerflow.runtime.runs.store.memory import MemoryRunStore
store = MemoryRunStore()
manager = RunManager(store=store)
old = await manager.create_or_reject("thread")
original = store.create_thread_operation_atomic
async def cancel_before_return(**kwargs):
result = await original(**kwargs)
asyncio.current_task().cancel()
return result
store.create_thread_operation_atomic = cancel_before_return
task = asyncio.create_task(manager.create_or_reject("thread", multitask_strategy="interrupt"))
with pytest.raises(asyncio.CancelledError):
await task
replacement = next(r for r in manager._runs.values() if r is not old)
assert (await store.get(replacement.run_id))["status"] == "interrupted"
@pytest.mark.anyio
async def test_delayed_atomic_interrupted_ack_cannot_restart_acknowledged_error():
from deerflow.runtime.runs.manager import TerminalCommitProof
from deerflow.runtime.runs.store.memory import MemoryRunStore
store = MemoryRunStore()
manager = manager_for(store)
record = await manager.create_or_reject("thread")
committed, release_ack = asyncio.Event(), asyncio.Event()
original = store.create_thread_operation_atomic
async def held_atomic(**kwargs):
result = await original(**kwargs)
committed.set()
await release_ack.wait()
return result
store.create_thread_operation_atomic = held_atomic
admission = asyncio.create_task(manager.create_or_reject("thread", multitask_strategy="rollback"))
try:
await asyncio.wait_for(committed.wait(), 1)
# A previously registered terminal writer can finish while admission
# awaits its older interrupted response under the manager lock.
outcome = TerminalCommitProof(RunStatus.error, "final error", None, None)
assert await manager._terminal_operation(record, outcome, lambda: store.update_status(record.run_id, "error", error="final error", owner_worker_id="owner"))
assert manager._has_final_terminal_ack(record)
release_ack.set()
await asyncio.wait_for(admission, 1)
assert record.status == RunStatus.error
assert record.error == "final error"
assert not record.abort_event.is_set()
assert not record.rollback_finalization_pending
assert not record.terminal_commit_pending
assert not manager._terminal_attempts
finally:
release_ack.set()
await asyncio.gather(admission, return_exceptions=True)
await manager.stop_heartbeat(timeout=0.1)

235
backend/tests/test_atomic_admission_worker.py Normal file
View File

@ -0,0 +1,235 @@
import asyncio
from datetime import UTC, datetime, timedelta
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import MessagesState, StateGraph
from test_terminal_store_contract import store as store
from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.runtime import RunManager, RunStatus
from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.runs.worker import RunContext, run_agent
@pytest.fixture
def anyio_backend():
return "asyncio"
@pytest.mark.anyio
@pytest.mark.parametrize(
"action,host_finalization_cancel,stale_renewal_race",
[
("interrupt", False, False),
("rollback", False, False),
("rollback", True, False),
("rollback", False, True),
],
)
@pytest.mark.parametrize("with_event_store", [True, False])
async def test_real_worker_atomic_replacement_retains_finalization_authority(store, action, with_event_store, host_finalization_cancel, stale_renewal_race):
manager = RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=5, grace_seconds=0))
record = await manager.create_or_reject("thread", user_id="review")
started, rollback_entered, release_rollback = (asyncio.Event() for _ in range(3))
error_committed, release_error_ack = asyncio.Event(), asyncio.Event()
hold_rollback = False
class Saver(InMemorySaver):
async def aput(self, config, checkpoint, metadata, new_versions):
if hold_rollback and metadata.get("source") == "update" and not rollback_entered.is_set():
rollback_entered.set()
await release_rollback.wait()
return await super().aput(config, checkpoint, metadata, new_versions)
saver = Saver()
async def blocked_node(state):
started.set()
await asyncio.Event().wait()
return {}
builder = StateGraph(MessagesState)
builder.add_node("step", blocked_node)
builder.set_entry_point("step")
builder.set_finish_point("step")
graph = builder.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "thread"}}
await graph.aupdate_state(config, {"messages": [HumanMessage(content="before")]}, as_node="step")
completion = AsyncMock()
bridge = SimpleNamespace(publish=AsyncMock(), publish_end=AsyncMock(), cleanup=AsyncMock())
record.task = asyncio.create_task(
run_agent(
bridge,
manager,
record,
ctx=RunContext(checkpointer=saver, event_store=MemoryRunEventStore() if with_event_store else None, on_run_completed=completion),
agent_factory=lambda **kwargs: graph,
graph_input={"messages": [HumanMessage(content="cancelled input")]},
config=config,
)
)
renewal = None
stale_renewal = None
stale_renewal_entered, release_stale_renewal = asyncio.Event(), asyncio.Event()
stale_renewal_calls = []
original_status = store.update_status
original_renew = store.renew_lease
async def hold_first_active_renewal(run_id, **kwargs):
stale_renewal_calls.append(kwargs.get("allow_interrupted", False))
if stale_renewal_race and len(stale_renewal_calls) == 1:
stale_renewal_entered.set()
await release_stale_renewal.wait()
return await original_renew(run_id, **kwargs)
store.renew_lease = hold_first_active_renewal
async def hold_final_ack(run_id, status, **kwargs):
result = await original_status(run_id, status, **kwargs)
if run_id == record.run_id and status == "error" and result is True:
error_committed.set()
await release_error_ack.wait()
return result
try:
await asyncio.wait_for(started.wait(), 2)
if stale_renewal_race:
stale_renewal = asyncio.create_task(manager._renew_leases())
await asyncio.wait_for(stale_renewal_entered.wait(), 2)
# Model the final failed-renewal window before replacement using an
# actual persisted deadline; the Worker and checkpoint code stay real.
deadline = datetime.now(UTC) + timedelta(milliseconds=150)
row = await store.get(record.run_id, user_id="review")
await store.put(record.run_id, thread_id="thread", status="running", user_id="review", owner_worker_id="owner", lease_expires_at=deadline.isoformat(), created_at=row["created_at"])
record.lease_expires_at = deadline.isoformat()
hold_rollback = action == "rollback"
if action == "rollback":
store.update_status = hold_final_ack
replacement = await manager.create_or_reject("thread", multitask_strategy=action, user_id="review")
assert replacement.status == RunStatus.pending
assert not record.ownership_lost
if action == "rollback":
await asyncio.wait_for(rollback_entered.wait(), 2)
assert not record.ownership_lost
assert record.terminal_commit_proof.status == RunStatus.interrupted
assert record.terminal_commit_proof.error == "Cancelled by newer run"
assert not manager.can_release_terminal_barrier(record)
assert not record.terminal_status_persisted
if stale_renewal_race:
release_stale_renewal.set()
await asyncio.wait_for(stale_renewal, 2)
assert record.status in (RunStatus.interrupted, RunStatus.error)
assert record.owner_worker_id == "owner"
assert manager._runs.get(record.run_id) is record
assert stale_renewal_calls == [False, True]
assert not record.ownership_lost
assert record.rollback_finalization_pending
await manager._renew_leases()
elapsed = asyncio.Event()
asyncio.get_running_loop().call_later(max(0, (deadline - datetime.now(UTC)).total_seconds()) + 0.03, elapsed.set)
await elapsed.wait()
row = await store.get(record.run_id, user_id="review")
assert datetime.now(UTC) > deadline
assert datetime.fromisoformat(row["lease_expires_at"]) > datetime.now(UTC)
assert row["status"] == "interrupted"
assert not record.task.done()
assert not record.ownership_lost
release_rollback.set()
await asyncio.wait_for(error_committed.wait(), 2)
assert not record.terminal_status_persisted
renewal = asyncio.create_task(manager._renew_leases())
await asyncio.sleep(0)
if host_finalization_cancel:
record.task.cancel("host cancellation during rollback terminal ACK")
await asyncio.sleep(0)
record.task.cancel("repeated host cancellation")
await asyncio.sleep(0)
assert not record.task.done()
assert not record.terminal_status_persisted
assert record.rollback_finalization_pending
assert not record.ownership_lost
release_error_ack.set()
if host_finalization_cancel:
with pytest.raises(asyncio.CancelledError, match="host cancellation during rollback terminal ACK"):
await asyncio.wait_for(record.task, 2)
else:
await asyncio.wait_for(record.task, 2)
if renewal:
await asyncio.wait_for(renewal, 1)
row = await store.get(record.run_id, user_id="review")
expected = "error" if action == "rollback" else "interrupted"
assert row["status"] == expected
assert record.status.value == expected
assert not record.ownership_lost
assert manager._has_terminal_ack(record)
assert record.run_id not in manager._lease_deadlines
completion.assert_awaited_once_with(record)
bridge.publish_end.assert_awaited_once_with(record.run_id)
if action == "rollback":
assert row["error"] == "Rolled back by user"
snapshot = await graph.aget_state(config)
assert [m.content for m in snapshot.values["messages"]] == ["before"]
finally:
release_rollback.set()
release_error_ack.set()
release_stale_renewal.set()
if not record.task.done():
record.task.cancel()
await asyncio.gather(record.task, *([renewal] if renewal else []), *([stale_renewal] if stale_renewal else []), return_exceptions=True)
await manager.stop_heartbeat(timeout=0.1)
@pytest.mark.anyio
async def test_fenced_worker_does_not_begin_rollback_checkpoint_writes(store):
manager = RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=5, grace_seconds=0))
record = await manager.create_or_reject("thread", user_id="review")
started = asyncio.Event()
saver = InMemorySaver()
async def blocked_node(state):
started.set()
await asyncio.Event().wait()
return {}
builder = StateGraph(MessagesState)
builder.add_node("step", blocked_node)
builder.set_entry_point("step")
builder.set_finish_point("step")
graph = builder.compile(checkpointer=saver)
config = {"configurable": {"thread_id": "thread"}}
await graph.aupdate_state(config, {"messages": [HumanMessage(content="before")]}, as_node="step")
completion = AsyncMock()
bridge = SimpleNamespace(publish=AsyncMock(), publish_end=AsyncMock(), cleanup=AsyncMock())
record.task = asyncio.create_task(
run_agent(
bridge,
manager,
record,
ctx=RunContext(checkpointer=saver, event_store=MemoryRunEventStore(), on_run_completed=completion),
agent_factory=lambda **kwargs: graph,
graph_input={"messages": [HumanMessage(content="cancelled input")]},
config=config,
)
)
try:
await asyncio.wait_for(started.wait(), 2)
before = await graph.aget_state(config)
# A rollback intent may already be local when its execution lease is
# lost. Cancellation delivered by fencing cannot grant rollback rights.
record.abort_action = "rollback"
record.lease_expires_at = (datetime.now(UTC) - timedelta(seconds=1)).isoformat()
await manager._renew_leases()
await asyncio.wait_for(record.task, 2)
assert record.ownership_lost
after = await graph.aget_state(config)
assert after.values == before.values
completion.assert_not_awaited()
finally:
if not record.task.done():
record.task.cancel()
await asyncio.gather(record.task, return_exceptions=True)
await manager.stop_heartbeat(timeout=0.1)

66
backend/tests/test_multi_worker_run_ownership.py
View File

@ -210,7 +210,7 @@ async def test_peer_idempotent_reuse_does_not_shield_crashed_owner_from_reconcil
# The owner crashes: its lease lapses past the grace window without renewal.
expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat()
assert await store.update_lease(first.run_id, owner_worker_id="worker-a", lease_expires_at=expired_lease)
store._runs[first.run_id]["lease_expires_at"] = expired_lease
recovered = await peer.reconcile_orphaned_inflight_runs(error="owner expired")
@ -454,8 +454,8 @@ async def test_reconciliation_skips_active_lease_runs():
@pytest.mark.anyio
async def test_reconciliation_skips_candidate_when_owner_renews_lease_after_scan():
"""A renewed lease between scan and claim must keep the run active."""
async def test_reconciliation_claims_when_expired_owner_attempts_renewal_after_scan():
"""An expired owner cannot revive its lease between scan and takeover."""
store = MemoryRunStore()
grace = 10
expired_lease = (datetime.now(UTC) - timedelta(seconds=grace + 5)).isoformat()
@ -477,7 +477,7 @@ async def test_reconciliation_skips_candidate_when_owner_renews_lease_after_scan
owner_worker_id="worker-alive",
lease_expires_at=renewed_lease,
)
assert updated is True
assert updated is False
return rows
store.list_inflight_with_expired_lease = list_then_owner_renews
@ -490,10 +490,10 @@ async def test_reconciliation_skips_candidate_when_owner_renews_lease_after_scan
error="Gateway restarted before this run reached a durable final state.",
)
assert recovered == []
assert [r.run_id for r in recovered] == ["race-run"]
stored = await store.get("race-run")
assert stored["status"] == "running"
assert datetime.fromisoformat(stored["lease_expires_at"]) > datetime.now(UTC)
assert stored["status"] == "error"
assert stored["lease_expires_at"] == expired_lease
@pytest.mark.anyio
@ -1072,7 +1072,6 @@ async def test_late_successful_renewal_still_fences_local_run():
near_expiry = (datetime.now(UTC) + timedelta(milliseconds=50)).isoformat()
record.lease_expires_at = near_expiry
store._runs[record.run_id]["lease_expires_at"] = near_expiry
original_update_lease = store.update_lease
async def renew_after_timeout_cancellation(*args, **kwargs):
try:
@ -1081,7 +1080,10 @@ async def test_late_successful_renewal_still_fences_local_run():
# Simulate a store operation that commits successfully despite the
# caller's deadline cancellation.
pass
return await original_update_lease(*args, **kwargs)
# Deliberately bypass the built-in Store's expiry guard to exercise
# the Manager against an unsafe/legacy Store's late positive result.
store._runs[args[0]]["lease_expires_at"] = kwargs["lease_expires_at"]
return True
store.update_lease = renew_after_timeout_cancellation
@ -1198,8 +1200,8 @@ async def test_renewal_rejected_by_own_terminal_commit_releases_barrier():
record.task = asyncio.create_task(asyncio.sleep(3600))
record.terminal_commit_pending = True
await manager.set_status(record.run_id, RunStatus.success, persist=False)
# The terminal write lands before the worker clears its barrier.
assert (await store.finalize_if_not_cancelled(record.run_id, status="success")).finalized
# Only a Manager-owned operation supplies an attributable acknowledgment.
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
try:
await manager._renew_leases()
@ -2556,7 +2558,7 @@ async def test_cancel_returns_taken_over_when_peer_claims_during_local_cancel():
# the race: in-memory cancel succeeds, but store write is blocked.
original = store.update_status
async def race_update(run_id, status, *, error=None, stop_reason=None):
async def race_update(run_id, status, *, error=None, stop_reason=None, **kwargs):
# Simulate peer takeover: flip to error before our write lands
run = store._runs.get(run_id)
if run and run["status"] == "running" and status == "interrupted":
@ -2564,7 +2566,7 @@ async def test_cancel_returns_taken_over_when_peer_claims_during_local_cancel():
run["error"] = "peer takeover"
run["updated_at"] = datetime.now(UTC).isoformat()
return False # our write was blocked
return await original(run_id, status, error=error, stop_reason=stop_reason)
return await original(run_id, status, error=error, stop_reason=stop_reason, **kwargs)
store.update_status = race_update
@ -2785,3 +2787,41 @@ def test_compute_retry_after_normal():
assert val is not None
# lease_expires_at is ~45s from now + grace_seconds 10 = ~55, within reason
assert 40 <= val <= 65
@pytest.mark.anyio
async def test_cancellation_during_terminal_receipt_cannot_commit_staged_success():
config = _lease_config(heartbeat_enabled=True)
store = MemoryRunStore()
manager = _make_manager(store=store, run_ownership_config=config)
record = await manager.create_or_reject("late-cancel")
entered, release = asyncio.Event(), asyncio.Event()
class ReceiptStore(MemoryRunEventStore):
async def put_if_absent(self, **event):
if event["event_type"] == "run.delivery":
entered.set()
await release.wait()
return await super().put_if_absent(**event)
class Agent:
async def astream(self, *args, **kwargs):
yield {"messages": []}
bridge = SimpleNamespace(publish=AsyncMock(), publish_end=AsyncMock(), cleanup=AsyncMock())
record.task = asyncio.create_task(run_agent(bridge, manager, record, ctx=RunContext(checkpointer=None, event_store=ReceiptStore()), agent_factory=lambda **kwargs: Agent(), graph_input={}, config={}))
try:
await asyncio.wait_for(entered.wait(), 1)
assert record.status == RunStatus.success
assert await store.request_cancel(record.run_id, action="interrupt") == "interrupt"
await manager._renew_leases()
assert record.abort_event.is_set()
release.set()
await asyncio.wait_for(record.task, 1)
assert not record.ownership_lost
assert record.status == RunStatus.interrupted
assert (await store.get(record.run_id))["status"] == "interrupted"
bridge.publish_end.assert_awaited_once()
finally:
release.set()
await asyncio.wait_for(asyncio.gather(record.task, return_exceptions=True), 1)

95
backend/tests/test_run_manager.py
View File

@ -1134,6 +1134,17 @@ async def test_create_or_reject_does_not_interrupt_old_run_when_new_run_store_wr
assert stored_old["status"] == "running"
def cancel_after_registration(monkeypatch, manager, old):
original_index = manager._index_run_locked
def index_and_cancel(record):
original_index(record)
if record is not old:
asyncio.get_running_loop().call_soon(asyncio.current_task().cancel)
monkeypatch.setattr(manager, "_index_run_locked", index_and_cancel)
@pytest.mark.anyio
@pytest.mark.parametrize("strategy", ["interrupt", "rollback"])
async def test_create_or_reject_cancellation_after_registration_interrupts_replacement(
@ -1145,20 +1156,8 @@ async def test_create_or_reject_cancellation_after_registration_interrupts_repla
manager = RunManager(store=store)
old = await manager.create("thread-1")
await manager.set_status(old.run_id, RunStatus.running)
persist_started = asyncio.Event()
release_persist = asyncio.Event()
original_persist_status = manager._persist_status
async def blocking_persist_status(record: Any, status: RunStatus, **kwargs: Any) -> bool:
persist_started.set()
await asyncio.wait_for(release_persist.wait(), timeout=1)
return await original_persist_status(record, status, **kwargs)
monkeypatch.setattr(manager, "_persist_status", blocking_persist_status)
cancel_after_registration(monkeypatch, manager, old)
create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy=strategy))
await asyncio.wait_for(persist_started.wait(), timeout=1)
create_task.cancel()
release_persist.set()
with pytest.raises(asyncio.CancelledError):
_ = await create_task
@ -1182,33 +1181,23 @@ async def test_create_or_reject_repeated_cancellation_drains_replacement_cleanup
manager = RunManager(store=store)
old = await manager.create("thread-1")
await manager.set_status(old.run_id, RunStatus.running)
old_persist_started = asyncio.Event()
release_old_persist = asyncio.Event()
replacement_persist_started = asyncio.Event()
release_replacement_persist = asyncio.Event()
original_persist_status = manager._persist_status
async def staged_persist_status(record: Any, status: RunStatus, **kwargs: Any) -> bool:
if record.run_id == old.run_id:
old_persist_started.set()
await asyncio.wait_for(release_old_persist.wait(), timeout=1)
else:
replacement_persist_started.set()
await asyncio.wait_for(release_replacement_persist.wait(), timeout=1)
replacement_persist_started.set()
await asyncio.wait_for(release_replacement_persist.wait(), timeout=1)
return await original_persist_status(record, status, **kwargs)
monkeypatch.setattr(manager, "_persist_status", staged_persist_status)
cancel_after_registration(monkeypatch, manager, old)
create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt"))
await asyncio.wait_for(old_persist_started.wait(), timeout=1)
replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id)
create_task.cancel()
release_old_persist.set()
await asyncio.wait_for(replacement_persist_started.wait(), timeout=1)
replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id)
create_task.cancel()
await asyncio.sleep(0)
assert not create_task.done()
release_replacement_persist.set()
with pytest.raises(asyncio.CancelledError):
_ = await asyncio.wait_for(create_task, timeout=1)
@ -1245,27 +1234,13 @@ async def test_create_or_reject_retries_replacement_when_cancel_status_cannot_pe
manager = RunManager(store=store)
old = await manager.create("thread-1", user_id="owner-1")
await manager.set_status(old.run_id, RunStatus.running)
old_persist_started = asyncio.Event()
release_old_persist = asyncio.Event()
original_persist_status = manager._persist_status
async def block_old_persist(record: Any, status: RunStatus, **kwargs: Any) -> bool:
if record.run_id != old.run_id:
return await original_persist_status(record, status, **kwargs)
old_persist_started.set()
await asyncio.wait_for(release_old_persist.wait(), timeout=1)
return await original_persist_status(record, status, **kwargs)
monkeypatch.setattr(manager, "_persist_status", block_old_persist)
cancel_after_registration(monkeypatch, manager, old)
create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt", user_id="owner-1"))
await asyncio.wait_for(old_persist_started.wait(), timeout=1)
replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id)
create_task.cancel()
release_old_persist.set()
with pytest.raises(asyncio.CancelledError):
_ = await asyncio.wait_for(create_task, timeout=1)
replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id)
stored_replacement = await store.get(replacement.run_id, user_id="owner-1")
assert stored_replacement is not None
assert stored_replacement["status"] == RunStatus.interrupted.value
@ -1299,22 +1274,8 @@ async def test_create_or_reject_cleanup_failure_preserves_caller_cancellation(
manager = RunManager(store=store)
old = await manager.create("thread-1")
await manager.set_status(old.run_id, RunStatus.running)
old_persist_started = asyncio.Event()
release_old_persist = asyncio.Event()
original_persist_status = manager._persist_status
async def block_old_persist(record: Any, status: RunStatus, **kwargs: Any) -> bool:
if record.run_id != old.run_id:
return await original_persist_status(record, status, **kwargs)
old_persist_started.set()
await asyncio.wait_for(release_old_persist.wait(), timeout=1)
return await original_persist_status(record, status, **kwargs)
monkeypatch.setattr(manager, "_persist_status", block_old_persist)
cancel_after_registration(monkeypatch, manager, old)
create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt"))
await asyncio.wait_for(old_persist_started.wait(), timeout=1)
create_task.cancel()
release_old_persist.set()
with pytest.raises(asyncio.CancelledError):
_ = await asyncio.wait_for(create_task, timeout=1)
@ -1343,27 +1304,13 @@ async def test_create_or_reject_preserves_peer_terminal_status_during_cancel_ret
manager = RunManager(store=store)
old = await manager.create("thread-1")
await manager.set_status(old.run_id, RunStatus.running)
old_persist_started = asyncio.Event()
release_old_persist = asyncio.Event()
original_persist_status = manager._persist_status
async def block_old_persist(record: Any, status: RunStatus, **kwargs: Any) -> bool:
if record.run_id != old.run_id:
return await original_persist_status(record, status, **kwargs)
old_persist_started.set()
await asyncio.wait_for(release_old_persist.wait(), timeout=1)
return await original_persist_status(record, status, **kwargs)
monkeypatch.setattr(manager, "_persist_status", block_old_persist)
cancel_after_registration(monkeypatch, manager, old)
create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt"))
await asyncio.wait_for(old_persist_started.wait(), timeout=1)
replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id)
create_task.cancel()
release_old_persist.set()
with pytest.raises(asyncio.CancelledError):
_ = await asyncio.wait_for(create_task, timeout=1)
replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id)
stored_replacement = await store.get(replacement.run_id)
assert stored_replacement is not None
assert stored_replacement["status"] == RunStatus.error.value

12
backend/tests/test_run_repository.py
View File

@ -1326,8 +1326,8 @@ class TestRunRepository:
await _cleanup()
@pytest.mark.anyio
async def test_reconciliation_skips_run_renewed_after_scan(self, tmp_path):
"""The SQL takeover CAS must reject a candidate renewed after its scan."""
async def test_reconciliation_claims_run_whose_expired_owner_attempted_renewal(self, tmp_path):
"""An expired lease cannot be renewed to evade the SQL takeover CAS."""
repo = await _make_repo(tmp_path)
grace = 10
run_id = "renewed-after-scan"
@ -1350,7 +1350,7 @@ class TestRunRepository:
owner_worker_id=owner_worker_id,
lease_expires_at=(datetime.now(UTC) + timedelta(seconds=60)).isoformat(),
)
assert renewed is True
assert renewed is False
return rows
repo.list_inflight_with_expired_lease = scan_then_renew
@ -1359,10 +1359,10 @@ class TestRunRepository:
recovered = await manager.reconcile_orphaned_inflight_runs(error="orphaned")
row = await repo.get(run_id)
assert recovered == []
assert [r.run_id for r in recovered] == [run_id]
assert row is not None
assert row["status"] == "running"
assert datetime.fromisoformat(row["lease_expires_at"]) > datetime.now(UTC)
assert row["status"] == "error"
assert datetime.fromisoformat(row["lease_expires_at"]) < datetime.now(UTC)
finally:
await _cleanup()

9
backend/tests/test_scheduled_goal_handoff.py
View File

@ -11,11 +11,13 @@ from langchain_core.messages import AIMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.graph import END, START, StateGraph
from sqlalchemy import update
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from deerflow.agents.thread_state import get_thread_state_schema
from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.persistence.base import Base
from deerflow.persistence.run.model import RunRow
from deerflow.persistence.run.sql import RunRepository
from deerflow.runtime import goal
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor
@ -397,7 +399,12 @@ async def test_expired_lease_fences_cleanup_until_recovery_reuses_the_thread(inf
async def evaluate(*args, **kwargs):
expiry = (datetime.now(UTC) - timedelta(seconds=1)).isoformat()
assert await store.update_lease(record.run_id, owner_worker_id="expired", lease_expires_at=expiry)
if isinstance(store, MemoryRunStore):
store._runs[record.run_id]["lease_expires_at"] = expiry
else:
async with store._sf() as session:
await session.execute(update(RunRow).where(RunRow.run_id == record.run_id).values(lease_expires_at=datetime.fromisoformat(expiry)))
await session.commit()
record.lease_expires_at = expiry
return {"satisfied": True, "blocker": "none", "reason": "Verified", "relied_on_assumption": False}

135
backend/tests/test_scheduled_goal_slow_rollback.py
View File

@ -1,9 +1,8 @@
"""Scheduled rollback keeps ownership until its saver and terminal CAS settle."""
import asyncio
import copy
import os
from datetime import UTC, datetime, timedelta
from datetime import UTC, datetime
from types import SimpleNamespace
from unittest.mock import AsyncMock
from uuid import uuid4
@ -212,10 +211,10 @@ async def test_terminal_commit_heartbeat_distinguishes_owner_outcome_from_takeov
error = "Rolled back by user" if status == "error" else None
await manager.set_status(record.run_id, RunStatus(status), error=error, goal_verdict=verdict, persist=False)
if writer == "owner":
outcome = await store.finalize_if_not_cancelled(record.run_id, status=status, error=error, goal_verdict=verdict)
assert outcome.finalized
await manager.set_status_if_not_cancelled(record.run_id, RunStatus(status), error=error, goal_verdict=verdict)
else:
assert await store.update_status(record.run_id, "error", error="Peer reclaimed this run", stop_reason="orphan_recovered")
# Identical fields and preserved owner still cannot prove attribution.
assert (await store.finalize_if_not_cancelled(record.run_id, status=status, error=error, goal_verdict=verdict)).finalized
# The durable terminal write has completed while the worker is still
# receiving its acknowledgement; the local finalization barrier is held.
await manager._renew_leases()
@ -224,77 +223,77 @@ async def test_terminal_commit_heartbeat_distinguishes_owner_outcome_from_takeov
@pytest.mark.asyncio
@pytest.mark.parametrize("status", ["success", "error"])
@pytest.mark.parametrize("path", ["renew_rejected", "expired_after_read"])
async def test_terminal_commit_after_detached_precheck_is_reconfirmed_before_fencing(database_runtime, monkeypatch, status, path):
@pytest.mark.parametrize("path", ["ack_before_renewal", "ack_during_renewal"])
async def test_terminal_attempt_ack_arbitrates_rejected_renewal(database_runtime, monkeypatch, status, path):
store, _ = database_runtime
manager = RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True))
record = await manager.create_or_reject("result", user_id="alice")
await manager.set_status(record.run_id, RunStatus.running)
record.scheduled_goal_cleanup_pending = True
verdict = {"satisfied": status == "success", "relied_on_assumption": False}
error = "Rolled back by user" if status == "error" else None
await manager.set_status(record.run_id, RunStatus(status), error=error, goal_verdict=verdict, persist=False)
get = store.get
intercepted = []
deadline = datetime.fromisoformat(record.lease_expires_at)
clock_advanced = False
committed, release, rejected = asyncio.Event(), asyncio.Event(), asyncio.Event()
finalize, renew = store.finalize_if_not_cancelled, store.renew_lease
if path == "expired_after_read":
from deerflow.runtime.runs import manager as manager_module
async def hold_ack(*args, **kwargs):
result = await finalize(*args, **kwargs)
committed.set()
await release.wait()
return result
class Clock(datetime):
@classmethod
def now(cls, tz=None):
return deadline + timedelta(seconds=1) if clock_advanced else datetime.now(tz)
async def observe_rejection(*args, **kwargs):
result = await renew(*args, **kwargs)
assert not result.renewed
rejected.set()
return result
monkeypatch.setattr(manager_module, "datetime", Clock)
async def commit_after_read(run_id, **kwargs):
nonlocal clock_advanced
snapshot = copy.deepcopy(await get(run_id, **kwargs))
if run_id == record.run_id and not intercepted:
intercepted.append(snapshot["status"])
result = await store.finalize_if_not_cancelled(run_id, status=status, error=error, goal_verdict=verdict)
assert result.finalized
clock_advanced = True
return snapshot
monkeypatch.setattr(store, "get", commit_after_read)
await manager._renew_leases()
assert intercepted == ["running"]
assert not record.ownership_lost
assert record.status.value == status
assert not record.scheduled_goal_cleanup_pending
monkeypatch.setattr(store, "finalize_if_not_cancelled", hold_ack)
monkeypatch.setattr(store, "renew_lease", observe_rejection)
commit = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus(status), error=error, goal_verdict=verdict))
renewal = None
try:
await asyncio.wait_for(committed.wait(), 2)
if path == "ack_before_renewal":
release.set()
await asyncio.wait_for(commit, 2)
renewal = asyncio.create_task(manager._renew_leases())
if path == "ack_during_renewal":
await asyncio.wait_for(rejected.wait(), 2)
assert record.scheduled_goal_cleanup_pending
assert not record.terminal_status_persisted
release.set()
await asyncio.wait_for(asyncio.gather(commit, renewal), 2)
assert not record.ownership_lost
assert record.status.value == status
assert not record.scheduled_goal_cleanup_pending
finally:
release.set()
await asyncio.wait_for(asyncio.gather(commit, *([renewal] if renewal else []), return_exceptions=True), 2)
@pytest.mark.asyncio
@pytest.mark.parametrize("status", ["success", "error"])
async def test_detached_precheck_commit_race_preserves_worker_hook_and_stream_end(database_runtime, monkeypatch, status):
async def test_delayed_terminal_ack_preserves_worker_hook_and_stream_end(database_runtime, monkeypatch, status):
store, checkpointer = database_runtime
manager = RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True))
metadata = {"scheduled_task_id": "task", "scheduled_task_run_id": "occurrence", "scheduled_goal_objective": "Report"}
record = await manager.create_or_reject("result", user_id="alice", metadata=metadata)
read_done, committed = asyncio.Event(), asyncio.Event()
get, finalize, update_status, write_goal = store.get, store.finalize_if_not_cancelled, store.update_status, worker.write_thread_goal
renewal_tasks = []
committed, release, rejected = asyncio.Event(), asyncio.Event(), asyncio.Event()
finalize, update_status, renew, write_goal = store.finalize_if_not_cancelled, store.update_status, store.renew_lease, worker.write_thread_goal
intercepted, cancelled = [], []
async def detached_precheck(run_id, **kwargs):
snapshot = copy.deepcopy(await get(run_id, **kwargs))
if renewal_tasks and asyncio.current_task() is renewal_tasks[0] and not intercepted:
intercepted.append(snapshot["status"])
read_done.set()
await committed.wait()
return snapshot
async def observe_rejection(*args, **kwargs):
result = await renew(*args, **kwargs)
if not result.renewed:
rejected.set()
return result
async def commit_with_delayed_ack(operation, run_id, **kwargs):
if run_id != record.run_id or kwargs["status"] != status or renewal_tasks:
return await operation(run_id, **kwargs)
renewal_tasks.append(asyncio.create_task(manager._renew_leases()))
await read_done.wait()
result = await operation(run_id, **kwargs)
committed.set()
await asyncio.shield(renewal_tasks[0])
if run_id == record.run_id and kwargs["status"] == status and not intercepted:
intercepted.append(status)
committed.set()
await release.wait()
return result
async def finalize_with_ack(run_id, **kwargs):
@ -315,7 +314,7 @@ async def test_detached_precheck_commit_race_preserves_worker_hook_and_stream_en
async def evaluate(*args, **kwargs):
return {"satisfied": True, "blocker": "none", "reason": "Verified", "relied_on_assumption": False}
monkeypatch.setattr(store, "get", detached_precheck)
monkeypatch.setattr(store, "renew_lease", observe_rejection)
monkeypatch.setattr(store, "finalize_if_not_cancelled", finalize_with_ack)
monkeypatch.setattr(store, "update_status", update_with_ack)
monkeypatch.setattr(worker, "write_thread_goal", cancel_for_error)
@ -334,12 +333,22 @@ async def test_detached_precheck_commit_race_preserves_worker_hook_and_stream_en
config={"configurable": {"thread_id": "result"}, "context": {"non_interactive": True}},
)
)
async with asyncio.timeout(10):
await asyncio.gather(record.task, return_exceptions=True)
if renewal_tasks:
await asyncio.gather(*renewal_tasks, return_exceptions=True)
assert intercepted == ["running"]
assert not record.ownership_lost
assert record.status.value == status
completed.assert_awaited_once_with(record)
bridge.publish_end.assert_awaited_once_with(record.run_id)
renewal = None
try:
await asyncio.wait_for(committed.wait(), 5)
renewal = asyncio.create_task(manager._renew_leases())
await asyncio.wait_for(rejected.wait(), 2)
assert not record.terminal_status_persisted
assert record.terminal_commit_pending
release.set()
await asyncio.wait_for(asyncio.gather(record.task, renewal), 5)
assert intercepted == [status]
assert not record.ownership_lost
assert record.status.value == status
completed.assert_awaited_once_with(record)
bridge.publish_end.assert_awaited_once_with(record.run_id)
finally:
release.set()
if not record.task.done():
record.task.cancel()
await asyncio.wait_for(asyncio.gather(record.task, *([renewal] if renewal else []), return_exceptions=True), 2)

177
backend/tests/test_stale_active_renewal.py Normal file
View File

@ -0,0 +1,177 @@
"""Lease renewals that predate rollback admission need a bounded second check."""
import asyncio
from contextlib import suppress
from datetime import UTC, datetime, timedelta
import pytest
from test_terminal_store_contract import store as store
from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.runtime import RunManager, RunStatus
def manager_for(run_store):
return RunManager(
store=run_store,
worker_id="owner",
run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=30, grace_seconds=0),
)
async def running_record(manager):
record = await manager.create_or_reject("thread")
await manager.set_status(record.run_id, RunStatus.running)
record.task = asyncio.create_task(asyncio.Event().wait())
return record
async def finish(manager, record, *tasks):
for task in tasks:
if not task.done():
task.cancel()
for task in tasks:
with suppress(asyncio.CancelledError):
await task
if not record.task.done():
record.task.cancel()
with suppress(asyncio.CancelledError):
await record.task
await manager.stop_heartbeat(timeout=0)
@pytest.mark.anyio
async def test_pre_rollback_renewal_retries_after_interrupted_ack(store):
manager = manager_for(store)
record = await running_record(manager)
entered, release = asyncio.Event(), asyncio.Event()
calls = []
original_renew = store.renew_lease
async def hold_first_active_renewal(run_id, **kwargs):
calls.append(kwargs.get("allow_interrupted", False))
if len(calls) == 1:
entered.set()
await release.wait()
return await original_renew(run_id, **kwargs)
store.renew_lease = hold_first_active_renewal
renewal = asyncio.create_task(manager._renew_leases())
try:
await asyncio.wait_for(entered.wait(), 1)
replacement = await manager.create_or_reject("thread", multitask_strategy="rollback")
assert replacement.status == RunStatus.pending
assert record.terminal_commit_proof is not None
assert record.terminal_commit_proof.status == RunStatus.interrupted
assert record.rollback_finalization_pending
old_expiry = record.lease_expires_at
release.set()
await asyncio.wait_for(renewal, 2)
assert calls == [False, True]
assert not record.ownership_lost
assert record.status == RunStatus.interrupted
assert record.rollback_finalization_pending
assert datetime.fromisoformat(record.lease_expires_at) > datetime.fromisoformat(old_expiry)
row = await store.get(record.run_id)
assert row["status"] == "interrupted"
assert row["owner_worker_id"] == "owner"
assert datetime.fromisoformat(row["lease_expires_at"]) > datetime.now(UTC)
finally:
release.set()
await finish(manager, record, renewal)
@pytest.mark.anyio
async def test_pre_rollback_renewal_waits_for_inflight_interrupted_ack(store):
manager = manager_for(store)
record = await running_record(manager)
renewal_entered, release_renewal = asyncio.Event(), asyncio.Event()
admission_committed, release_admission_ack = asyncio.Event(), asyncio.Event()
calls = []
original_renew = store.renew_lease
original_admission = store.create_thread_operation_atomic
async def hold_first_active_renewal(run_id, **kwargs):
calls.append(kwargs.get("allow_interrupted", False))
if len(calls) == 1:
renewal_entered.set()
await release_renewal.wait()
return await original_renew(run_id, **kwargs)
async def hold_committed_admission_ack(*args, **kwargs):
result = await original_admission(*args, **kwargs)
admission_committed.set()
await release_admission_ack.wait()
return result
store.renew_lease = hold_first_active_renewal
store.create_thread_operation_atomic = hold_committed_admission_ack
renewal = asyncio.create_task(manager._renew_leases())
admission = None
try:
await asyncio.wait_for(renewal_entered.wait(), 1)
admission = asyncio.create_task(manager.create_or_reject("thread", multitask_strategy="rollback"))
await asyncio.wait_for(admission_committed.wait(), 1)
assert record.terminal_commit_proof is None
release_renewal.set()
await asyncio.sleep(0)
assert not renewal.done()
assert not record.ownership_lost
release_admission_ack.set()
replacement = await asyncio.wait_for(admission, 2)
await asyncio.wait_for(renewal, 2)
assert replacement.status == RunStatus.pending
assert calls == [False, True]
assert record.terminal_commit_proof is not None
assert record.terminal_commit_proof.status == RunStatus.interrupted
assert not record.ownership_lost
assert (await store.get(record.run_id))["status"] == "interrupted"
finally:
release_renewal.set()
release_admission_ack.set()
await finish(manager, record, renewal, *([admission] if admission is not None else []))
@pytest.mark.anyio
@pytest.mark.parametrize("taken_over", [False, True], ids=["expired", "peer"])
async def test_pre_rollback_retry_still_fences_when_live_owner_was_lost(store, taken_over):
manager = manager_for(store)
record = await running_record(manager)
entered, release = asyncio.Event(), asyncio.Event()
calls = []
original_renew = store.renew_lease
async def lose_owner_before_retry(run_id, **kwargs):
calls.append(kwargs.get("allow_interrupted", False))
if len(calls) == 1:
entered.set()
await release.wait()
return await original_renew(run_id, **kwargs)
store.renew_lease = lose_owner_before_retry
renewal = asyncio.create_task(manager._renew_leases())
try:
await asyncio.wait_for(entered.wait(), 1)
await manager.create_or_reject("thread", multitask_strategy="rollback")
await store.put(
record.run_id,
thread_id=record.thread_id,
status="interrupted",
owner_worker_id="peer" if taken_over else "owner",
lease_expires_at=(datetime.now(UTC) + (timedelta(seconds=30) if taken_over else -timedelta(seconds=1))).isoformat(),
error="Cancelled by newer run",
)
release.set()
await asyncio.wait_for(renewal, 2)
assert calls == [False, True]
assert record.ownership_lost
finally:
release.set()
await finish(manager, record, renewal)

576
backend/tests/test_terminal_commit_authority.py Normal file
View File

@ -0,0 +1,576 @@
"""Deterministic terminal attribution and independent lease deadline regressions."""
import asyncio
from contextlib import suppress
from datetime import UTC, datetime, timedelta
from unittest.mock import AsyncMock
import pytest
from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.runtime import RunManager, RunStatus
from deerflow.runtime.runs.store.base import RunStore, StatusFinalization
from deerflow.runtime.runs.store.memory import MemoryRunStore
def manager_for(store):
return RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=30, grace_seconds=0))
async def staged(manager, thread="thread"):
record = await manager.create_or_reject(thread)
record.terminal_commit_pending = True
await manager.set_status(record.run_id, RunStatus.success, persist=False)
return record
async def finish(*tasks):
for task in tasks:
if not task.done():
task.cancel()
for task in tasks:
with suppress(asyncio.CancelledError):
await asyncio.wait_for(task, 1)
@pytest.mark.anyio
@pytest.mark.parametrize("status", [RunStatus.success, RunStatus.error])
async def test_identical_peer_terminal_row_is_not_own_ack(status):
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
await manager.set_status(record.run_id, status, persist=False)
assert (await store.finalize_if_not_cancelled(record.run_id, status=status.value)).finalized
await manager._renew_leases()
assert record.ownership_lost, "an externally committed identical row must fence the unconfirmed writer"
@pytest.mark.anyio
async def test_rejected_cas_has_no_status_or_completion_fallback():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
store.finalize_if_not_cancelled = AsyncMock(return_value=StatusFinalization(finalized=False))
store.update_status = AsyncMock(wraps=store.update_status)
store.update_run_completion = AsyncMock(wraps=store.update_run_completion)
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
await manager.update_run_completion(record.run_id, status="success", total_tokens=7)
assert (await store.get(record.run_id))["status"] == "pending"
store.update_status.assert_not_awaited()
store.update_run_completion.assert_not_awaited()
assert record.ownership_lost
@pytest.mark.anyio
async def test_unacknowledged_completion_cannot_terminalize_or_recreate():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
await manager.update_run_completion(record.run_id, status="success", total_tokens=7)
assert (await store.get(record.run_id))["status"] == "pending"
assert not record.terminal_status_persisted
await store.delete(record.run_id)
await manager.update_run_completion(record.run_id, status="success")
assert await store.get(record.run_id) is None
@pytest.mark.anyio
async def test_committed_cas_waits_for_its_ack_after_renewal_rejection():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
committed, release, rejected = asyncio.Event(), asyncio.Event(), asyncio.Event()
original_finalize, original_renew = store.finalize_if_not_cancelled, store.renew_lease
async def commit_then_hold_ack(*args, **kwargs):
result = await original_finalize(*args, **kwargs)
committed.set()
await release.wait()
return result
async def reject_terminal_renewal(*args, **kwargs):
result = await original_renew(*args, **kwargs)
rejected.set()
return result
store.finalize_if_not_cancelled = commit_then_hold_ack
store.renew_lease = reject_terminal_renewal
commit = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success))
renewal = None
try:
await asyncio.wait_for(committed.wait(), 1)
renewal = asyncio.create_task(manager._renew_leases())
await asyncio.wait_for(rejected.wait(), 0.3)
assert record.terminal_commit_pending
assert not renewal.done(), "the matching durable row cannot replace the held acknowledgment"
release.set()
await asyncio.wait_for(asyncio.gather(commit, renewal), 1)
assert not record.ownership_lost
assert not record.terminal_commit_pending
assert record.terminal_status_persisted
finally:
release.set()
await finish(commit, *([renewal] if renewal else []))
@pytest.mark.anyio
async def test_uncancellable_renewal_does_not_delay_peer_cancel_or_deadline():
store = MemoryRunStore()
manager = manager_for(store)
slow = await manager.create_or_reject("slow")
fast = await manager.create_or_reject("fast")
deadline = (datetime.now(UTC) + timedelta(milliseconds=100)).isoformat()
slow.lease_expires_at = deadline
store._runs[slow.run_id]["lease_expires_at"] = deadline
entered, release, fast_renewed = asyncio.Event(), asyncio.Event(), asyncio.Event()
original = store.renew_lease
calls = 0
async def blocked(*args, **kwargs):
nonlocal calls
if args[0] == slow.run_id:
calls += 1
entered.set()
while not release.is_set():
try:
await release.wait()
except asyncio.CancelledError:
pass
return await original(*args, **kwargs)
result = await original(*args, **kwargs)
fast_renewed.set()
return result
store.renew_lease = blocked
await store.request_cancel(fast.run_id, action="rollback")
cycle = asyncio.create_task(manager._renew_leases())
try:
await asyncio.wait_for(entered.wait(), 1)
await asyncio.wait_for(fast_renewed.wait(), 0.3)
await asyncio.wait_for(asyncio.shield(cycle), 0.4)
assert slow.ownership_lost
assert fast.abort_event.is_set()
assert fast.abort_action == "rollback"
assert not fast.ownership_lost
await manager._renew_leases()
assert calls == 1
release.set()
await manager.stop_heartbeat(timeout=0.1)
assert slow.ownership_lost
assert slow.lease_expires_at == deadline
finally:
release.set()
await finish(cycle)
await manager.stop_heartbeat(timeout=0.1)
@pytest.mark.anyio
async def test_caller_cancelled_cas_has_no_positive_proof():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
entered, release = asyncio.Event(), asyncio.Event()
original = store.finalize_if_not_cancelled
async def commit_without_ack(*args, **kwargs):
await original(*args, **kwargs)
entered.set()
await release.wait()
return StatusFinalization(finalized=True)
store.finalize_if_not_cancelled = commit_without_ack
commit = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success))
try:
await asyncio.wait_for(entered.wait(), 1)
commit.cancel()
with pytest.raises(asyncio.CancelledError):
await commit
await manager._renew_leases()
assert record.ownership_lost
assert not record.terminal_status_persisted
finally:
release.set()
await finish(commit)
@pytest.mark.anyio
async def test_different_attempts_settle_without_borrowing_another_status_ack():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
committed, rejected, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
original = store.finalize_if_not_cancelled
async def hold_success(*args, **kwargs):
result = await original(*args, **kwargs)
if result.finalized:
committed.set()
await release.wait()
else:
rejected.set()
return result
store.finalize_if_not_cancelled = hold_success
success = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success))
failure = None
try:
await asyncio.wait_for(committed.wait(), 1)
failure = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.error, error="losing attempt"))
await asyncio.wait_for(rejected.wait(), 1)
assert record.terminal_commit_pending
assert not record.terminal_status_persisted
release.set()
await asyncio.wait_for(asyncio.gather(success, failure), 1)
assert record.status == RunStatus.success
assert not record.ownership_lost
assert record.terminal_status_persisted
assert manager._terminal_attempts == {}
finally:
release.set()
await finish(success, *([failure] if failure else []))
@pytest.mark.anyio
async def test_unknown_commit_fences_at_deadline_and_late_ack_cannot_revive():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
expiry = (datetime.now(UTC) + timedelta(milliseconds=100)).isoformat()
record.lease_expires_at = expiry
store._runs[record.run_id]["lease_expires_at"] = expiry
committed, release = asyncio.Event(), asyncio.Event()
original = store.finalize_if_not_cancelled
async def hold_ack(*args, **kwargs):
result = await original(*args, **kwargs)
committed.set()
await release.wait()
return result
store.finalize_if_not_cancelled = hold_ack
commit = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success))
try:
await asyncio.wait_for(committed.wait(), 1)
await asyncio.wait_for(manager._renew_leases(), 0.5)
assert record.ownership_lost
assert not record.terminal_status_persisted
release.set()
await asyncio.wait_for(commit, 1)
assert record.ownership_lost
assert record.status == RunStatus.error
assert not record.terminal_status_persisted
assert manager._terminal_attempts == {}
finally:
release.set()
await finish(commit)
@pytest.mark.anyio
async def test_legacy_none_status_result_is_unconfirmed_in_heartbeat_mode():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
original = store.update_status
async def no_rowcount(*args, **kwargs):
await original(*args, **kwargs)
store.update_status = no_rowcount
assert await manager.persist_current_status(record.run_id) is False
assert record.ownership_lost
assert not record.terminal_status_persisted
@pytest.mark.anyio
@pytest.mark.parametrize("method", ["update_status", "legacy_finalization", "default_finalization"])
async def test_unsupported_owned_store_contract_fails_closed_with_diagnostic(method, caplog):
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
calls = []
async def legacy_write(run_id, status, *, error=None, stop_reason=None, goal_verdict=None):
calls.append(run_id)
return True
if method == "update_status":
store.update_status = legacy_write
assert await manager.persist_current_status(record.run_id) is False
else:
store.finalize_if_not_cancelled = legacy_write if method == "legacy_finalization" else RunStore.finalize_if_not_cancelled.__get__(store)
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
assert calls == []
assert (await store.get(record.run_id))["status"] == "pending"
assert record.ownership_lost
assert record.terminal_commit_proof is None
assert not record.terminal_status_persisted
assert manager._terminal_attempts == {}
assert "owner_worker_id" in caplog.text
@pytest.mark.anyio
async def test_legacy_completion_cannot_bypass_strict_contract(caplog):
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
proof = record.terminal_commit_proof
assert proof is not None
calls = []
async def legacy_completion(run_id, *, status, total_tokens=0):
calls.append(run_id)
return True
store.update_run_completion = legacy_completion
await manager.update_run_completion(record.run_id, status="success", total_tokens=5)
assert calls == []
assert (await store.get(record.run_id)).get("total_tokens", 0) == 0
assert record.terminal_commit_proof is proof
assert not record.ownership_lost
assert "require_terminal" in caplog.text
@pytest.mark.anyio
async def test_subsequent_ticks_renew_fast_run_while_slow_check_is_pending():
store = MemoryRunStore()
manager = manager_for(store)
slow = await manager.create_or_reject("slow")
fast = await manager.create_or_reject("fast")
entered, release, twice = asyncio.Event(), asyncio.Event(), asyncio.Event()
original = store.renew_lease
slow_calls, fast_calls = 0, 0
async def renew(run_id, **kwargs):
nonlocal slow_calls, fast_calls
if run_id == slow.run_id:
slow_calls += 1
entered.set()
await release.wait()
result = await original(run_id, **kwargs)
if run_id == fast.run_id:
fast_calls += 1
if fast_calls == 2:
twice.set()
return result
store.renew_lease = renew
try:
await manager._renew_leases(wait=False)
await asyncio.wait_for(entered.wait(), 1)
async with asyncio.timeout(1):
while fast_calls < 1:
await asyncio.sleep(0)
await asyncio.wait_for(asyncio.shield(manager._lease_checks[fast.run_id]), 1)
await manager._renew_leases(wait=False)
await asyncio.wait_for(twice.wait(), 1)
assert slow_calls == 1
assert not manager._lease_checks[slow.run_id].done()
assert not fast.ownership_lost
finally:
release.set()
await manager.stop_heartbeat(timeout=1)
@pytest.mark.anyio
async def test_shutdown_bounds_uncancellable_io_and_observes_late_exception(caplog):
store = MemoryRunStore()
manager = manager_for(store)
record = await manager.create_or_reject("thread")
entered, release = asyncio.Event(), asyncio.Event()
async def resistant(*args, **kwargs):
entered.set()
while not release.is_set():
try:
await release.wait()
except asyncio.CancelledError:
pass
raise RuntimeError("late Store failure")
store.renew_lease = resistant
try:
await manager._renew_leases(wait=False)
await asyncio.wait_for(entered.wait(), 1)
await asyncio.wait_for(manager.stop_heartbeat(timeout=0.01), 0.3)
assert record.run_id in manager._lease_operations
assert "still draining" in caplog.text
release.set()
async with asyncio.timeout(1):
while manager._lease_operations:
await asyncio.sleep(0)
assert "RuntimeError" in caplog.text
finally:
release.set()
await manager.stop_heartbeat(timeout=1)
@pytest.mark.anyio
async def test_staged_terminal_still_receives_durable_cancellation():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
assert await store.request_cancel(record.run_id, action="rollback") == "rollback"
await manager._renew_leases()
assert record.abort_event.is_set()
assert record.abort_action == "rollback"
assert not record.ownership_lost
@pytest.mark.anyio
async def test_completion_cannot_mutate_an_acknowledged_outcome():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
await manager.update_run_completion(record.run_id, status="success", error="different outcome")
assert record.error is None
assert (await store.get(record.run_id)).get("error") is None
assert manager._has_terminal_ack(record)
@pytest.mark.anyio
async def test_exited_unconfirmed_worker_does_not_keep_renewing_forever():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
record.task = asyncio.create_task(asyncio.sleep(0))
await record.task
await manager._renew_leases()
assert record.ownership_lost
assert not record.terminal_commit_pending
@pytest.mark.anyio
async def test_barrier_requires_own_ack_only_for_durable_heartbeat_mode():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
assert not manager.can_release_terminal_barrier(record)
memory_only = RunManager(run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True))
assert memory_only.can_release_terminal_barrier(record)
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
assert manager.can_release_terminal_barrier(record)
@pytest.mark.anyio
async def test_older_interrupted_ack_cannot_overwrite_newer_error_commit():
store = MemoryRunStore()
manager = manager_for(store)
record = await manager.create_or_reject("thread")
committed, release = asyncio.Event(), asyncio.Event()
original = store.update_status
async def hold_interrupted(*args, **kwargs):
result = await original(*args, **kwargs)
if args[1] == "interrupted":
committed.set()
await release.wait()
return result
store.update_status = hold_interrupted
first = asyncio.create_task(manager.set_status(record.run_id, RunStatus.interrupted))
try:
await asyncio.wait_for(committed.wait(), 1)
await manager.set_status(record.run_id, RunStatus.error, error="Rolled back by user")
assert record.status == RunStatus.error
release.set()
await asyncio.wait_for(first, 1)
assert record.status == RunStatus.error
assert (await store.get(record.run_id))["status"] == "error"
assert record.error == "Rolled back by user"
assert record.terminal_status_persisted
assert manager._has_terminal_ack(record)
assert not manager._has_terminal_ack(record, manager._terminal_outcome(record, RunStatus.interrupted, error="Rolled back by user"))
assert not record.ownership_lost
assert manager._terminal_attempts == {}
finally:
release.set()
await finish(first)
@pytest.mark.anyio
async def test_attempt_uses_immutable_verdict_snapshot_for_store_and_proof():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
entered, release = asyncio.Event(), asyncio.Event()
original = store.finalize_if_not_cancelled
verdict = {"satisfied": True}
async def delayed(*args, **kwargs):
entered.set()
await release.wait()
return await original(*args, **kwargs)
store.finalize_if_not_cancelled = delayed
commit = asyncio.create_task(manager.set_status_if_not_cancelled(record.run_id, RunStatus.success, goal_verdict=verdict))
try:
await asyncio.wait_for(entered.wait(), 1)
verdict["satisfied"] = False
release.set()
await asyncio.wait_for(commit, 1)
assert (await store.get(record.run_id))["goal_verdict"] == {"satisfied": True}
assert record.goal_verdict == {"satisfied": True}
assert manager._has_terminal_ack(record)
finally:
release.set()
await finish(commit)
@pytest.mark.anyio
async def test_intermediate_ack_cannot_confirm_newer_pending_error():
store = MemoryRunStore()
manager = manager_for(store)
record = await manager.create_or_reject("thread")
first_committed, second_committed = asyncio.Event(), asyncio.Event()
release_first, release_second = asyncio.Event(), asyncio.Event()
original = store.update_status
async def hold_ack(run_id, status, **kwargs):
result = await original(run_id, status, **kwargs)
if status == "interrupted":
first_committed.set()
await release_first.wait()
elif status == "error":
second_committed.set()
await release_second.wait()
return result
store.update_status = hold_ack
first = asyncio.create_task(manager.set_status(record.run_id, RunStatus.interrupted))
second = None
try:
await asyncio.wait_for(first_committed.wait(), 1)
second = asyncio.create_task(manager.set_status(record.run_id, RunStatus.error, error="Rolled back by user"))
await asyncio.wait_for(second_committed.wait(), 1)
release_first.set()
await asyncio.wait_for(first, 1)
assert record.status == RunStatus.error
assert not record.terminal_status_persisted
assert record.terminal_commit_pending
release_second.set()
await asyncio.wait_for(second, 1)
assert record.status == RunStatus.error
assert record.terminal_status_persisted
assert not record.terminal_commit_pending
assert not record.ownership_lost
finally:
release_first.set()
release_second.set()
await finish(first, *([second] if second else []))
@pytest.mark.anyio
async def test_acknowledged_final_outcome_rejects_late_conflicting_local_status():
store = MemoryRunStore()
manager = manager_for(store)
record = await staged(manager)
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
await manager.set_status(record.run_id, RunStatus.error, error="late cancellation")
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.error, error="late failure")
assert record.status == RunStatus.success
assert (await store.get(record.run_id))["status"] == "success"
assert manager._has_terminal_ack(record)
assert not record.ownership_lost

93
backend/tests/test_terminal_completion_capability.py Normal file
View File

@ -0,0 +1,93 @@
"""Third-party completion stores must opt into the strict metadata contract."""
import pytest
from deerflow.config.run_ownership_config import RunOwnershipConfig
from deerflow.runtime import RunManager, RunStatus
from deerflow.runtime.runs.store.memory import MemoryRunStore
class AbsorbingStore(MemoryRunStore):
async def update_run_completion(self, run_id, *, status, **kwargs):
self.completion_calls = getattr(self, "completion_calls", 0) + 1
self._runs[run_id].update(status=status, **kwargs)
return True
class FlagOnlyChild(AbsorbingStore):
supports_terminal_completion_kwargs = True
class OptedInStore(MemoryRunStore):
supports_terminal_completion_kwargs = True
async def update_run_completion(self, run_id, *, status, **kwargs):
self.completion_calls = getattr(self, "completion_calls", 0) + 1
return await super().update_run_completion(run_id, status=status, **kwargs)
class InheritedOptIn(OptedInStore):
pass
class UnoptedOverride(OptedInStore):
async def update_run_completion(self, run_id, *, status, **kwargs):
self.completion_calls = getattr(self, "completion_calls", 0) + 1
return await MemoryRunStore.update_run_completion(self, run_id, status=status, **kwargs)
class ExplicitStore(MemoryRunStore):
async def update_run_completion(self, run_id, *, status, require_terminal=False, owner_worker_id=None, **kwargs):
self.completion_calls = getattr(self, "completion_calls", 0) + 1
return await super().update_run_completion(run_id, status=status, require_terminal=require_terminal, owner_worker_id=owner_worker_id, **kwargs)
class PartialContract(AbsorbingStore):
async def update_run_completion(self, run_id, *, status, require_terminal=False, **kwargs):
return await super().update_run_completion(run_id, status=status, require_terminal=require_terminal, **kwargs)
@pytest.mark.anyio
@pytest.mark.parametrize("store_type", [AbsorbingStore, FlagOnlyChild, UnoptedOverride, PartialContract])
async def test_unproven_kwargs_completion_is_skipped(store_type, caplog):
store = store_type()
manager = RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True))
record = await manager.create_or_reject("thread")
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
proof = record.terminal_commit_proof
assert proof is not None
await manager.update_run_completion(record.run_id, status="success", total_tokens=7)
assert getattr(store, "completion_calls", 0) == 0
row = await store.get(record.run_id)
assert row.get("total_tokens", 0) == 0
assert "require_terminal" not in row
assert record.terminal_commit_proof is proof
assert not record.ownership_lost
assert "require_terminal" in caplog.text
assert "owner_worker_id" in caplog.text
@pytest.mark.anyio
@pytest.mark.parametrize("store_type", [MemoryRunStore, ExplicitStore, OptedInStore, InheritedOptIn])
async def test_declared_completion_contract_updates_metadata(store_type):
store = store_type()
manager = RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True))
record = await manager.create_or_reject("thread")
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
await manager.update_run_completion(record.run_id, status="success", total_tokens=7)
row = await store.get(record.run_id)
assert row["total_tokens"] == 7
assert "require_terminal" not in row
assert row["owner_worker_id"] == "owner"
@pytest.mark.anyio
async def test_single_worker_preserves_legacy_kwargs_completion():
store = AbsorbingStore()
manager = RunManager(store=store)
record = await manager.create_or_reject("thread")
await manager.update_run_completion(record.run_id, status="success", total_tokens=7)
assert store.completion_calls == 1
row = await store.get(record.run_id)
assert row["total_tokens"] == 7
assert "require_terminal" not in row

196
backend/tests/test_terminal_store_contract.py Normal file
View File

@ -0,0 +1,196 @@
"""Execute the terminal/lease contract against memory and real SQL backends."""
import os
import uuid
from datetime import UTC, datetime, timedelta
import pytest
from sqlalchemy import text
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from support.postgres import asyncpg_test_url
from deerflow.persistence.run import RunRepository
from deerflow.persistence.run.model import RunChangeClockRow, RunRow
from deerflow.runtime.runs.store.memory import MemoryRunStore
@pytest.fixture(params=["memory", "sqlite", "postgresql"])
async def store(request, tmp_path):
if request.param == "memory":
yield MemoryRunStore()
return
schema = None
if request.param == "postgresql":
uri = os.getenv("TEST_POSTGRES_URI")
if not uri:
pytest.skip("set TEST_POSTGRES_URI for real PostgreSQL terminal/lease contracts")
schema = f"r2_{uuid.uuid4().hex}"
engine = create_async_engine(asyncpg_test_url(uri), connect_args={"server_settings": {"search_path": schema}})
else:
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'r2.db'}")
try:
async with engine.begin() as conn:
if schema:
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
await conn.run_sync(lambda c: RunRow.metadata.create_all(c, tables=[RunRow.__table__, RunChangeClockRow.__table__]))
yield RunRepository(async_sessionmaker(engine, expire_on_commit=False))
finally:
if schema:
async with engine.begin() as conn:
await conn.execute(text(f'DROP SCHEMA "{schema}" CASCADE'))
await engine.dispose()
async def put(store, *, expired=False):
expiry = (datetime.now(UTC) + timedelta(seconds=-1 if expired else 30)).isoformat()
await store.put("run", thread_id="thread", status="running", owner_worker_id="owner", lease_expires_at=expiry)
return expiry
@pytest.mark.anyio
async def test_expired_lease_cannot_be_resurrected(store):
expired = await put(store, expired=True)
renewal = await store.renew_lease("run", owner_worker_id="owner", lease_expires_at=(datetime.now(UTC) + timedelta(seconds=60)).isoformat())
assert not renewal.renewed
assert datetime.fromisoformat((await store.get("run"))["lease_expires_at"]) == datetime.fromisoformat(expired)
@pytest.mark.anyio
@pytest.mark.parametrize("expired,owner", [(True, "owner"), (False, "peer")])
async def test_terminal_cas_requires_live_owner(store, expired, owner):
await put(store, expired=expired)
result = await store.finalize_if_not_cancelled("run", status="success", owner_worker_id=owner)
assert not result.finalized
assert (await store.get("run"))["status"] == "running"
@pytest.mark.anyio
async def test_cancel_wins_and_metadata_cannot_bypass_cas(store):
await put(store)
assert await store.request_cancel("run", action="rollback") == "rollback"
result = await store.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner")
assert not result.finalized
assert result.cancel_action == "rollback"
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="owner", total_tokens=5) is False
assert (await store.get("run"))["status"] == "running"
@pytest.mark.anyio
async def test_confirmed_terminal_allows_only_matching_metadata(store):
await put(store)
assert (await store.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner")).finalized
assert await store.update_run_completion("run", status="error", require_terminal=True, owner_worker_id="owner") is False
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="peer") is False
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="owner", total_tokens=5) is True
assert (await store.get("run"))["total_tokens"] == 5
@pytest.mark.anyio
async def test_terminal_metadata_needs_matching_owner_and_outcome_but_not_live_lease(store):
expired = (datetime.now(UTC) - timedelta(seconds=1)).isoformat()
await store.put("run", thread_id="thread", status="success", owner_worker_id="owner", lease_expires_at=expired)
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="owner", total_tokens=5)
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="peer", total_tokens=10) is False
assert await store.update_run_completion("run", status="error", require_terminal=True, owner_worker_id="owner", total_tokens=10) is False
row = await store.get("run")
assert row["status"] == "success"
assert row["total_tokens"] == 5
assert datetime.fromisoformat(row["lease_expires_at"]) == datetime.fromisoformat(expired)
@pytest.mark.anyio
@pytest.mark.parametrize("method", ["renew_lease", "update_lease", "finalize_if_not_cancelled", "update_status"])
async def test_postgres_checks_database_time_after_row_lock(store, monkeypatch, method):
"""A transaction begun before expiry cannot write after waiting past it."""
if not isinstance(store, RunRepository) or store._sf.kw["bind"].dialect.name != "postgresql":
pytest.skip("requires PostgreSQL row locks and clock_timestamp")
import asyncio
from sqlalchemy import func, select
expiry = datetime.now(UTC) + timedelta(milliseconds=300)
await store.put("locked", thread_id="locked-thread", status="running", owner_worker_id="owner", lease_expires_at=expiry.isoformat())
entered = asyncio.Event()
lock = store._lock_owned_row
async def observe_transaction_start(session, run_id):
assert await session.scalar(select(func.clock_timestamp())) < expiry
entered.set()
await lock(session, run_id)
monkeypatch.setattr(store, "_lock_owned_row", observe_transaction_start)
operation = None
async with store._sf() as blocker:
await blocker.execute(select(RunRow).where(RunRow.run_id == "locked").with_for_update())
try:
if method in ("renew_lease", "update_lease"):
operation = asyncio.create_task(getattr(store, method)("locked", owner_worker_id="owner", lease_expires_at=(expiry + timedelta(seconds=30)).isoformat()))
elif method == "update_status":
operation = asyncio.create_task(store.update_status("locked", "error", owner_worker_id="owner"))
else:
operation = asyncio.create_task(store.finalize_if_not_cancelled("locked", status="success", owner_worker_id="owner"))
await asyncio.wait_for(entered.wait(), 1)
async with asyncio.timeout(2):
while await blocker.scalar(select(func.clock_timestamp())) <= expiry:
await asyncio.sleep(0.01)
assert not operation.done()
finally:
await blocker.rollback()
assert operation is not None
result = await asyncio.wait_for(operation, 1)
if method == "renew_lease":
assert not result.renewed
elif method == "finalize_if_not_cancelled":
assert not result.finalized
else:
assert result is False
row = await store.get("locked")
assert row["status"] == "running"
assert datetime.fromisoformat(row["lease_expires_at"]) == expiry
@pytest.mark.anyio
async def test_durable_cancel_does_not_authorize_expired_terminal_writer(store):
await put(store, expired=True)
assert await store.request_cancel("run", action="rollback") == "rollback"
result = await store.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner")
assert not result.finalized
assert result.cancel_action is None
@pytest.mark.anyio
async def test_renewal_cannot_write_an_already_expired_target(store):
original = await put(store)
renewal = await store.renew_lease("run", owner_worker_id="owner", lease_expires_at=(datetime.now(UTC) - timedelta(seconds=1)).isoformat())
assert not renewal.renewed
assert datetime.fromisoformat((await store.get("run"))["lease_expires_at"]) == datetime.fromisoformat(original)
@pytest.mark.anyio
async def test_strict_completion_cannot_rewrite_terminal_outcome_fields(store):
await put(store)
verdict = {"satisfied": True}
assert (await store.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner", goal_verdict=verdict)).finalized
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="owner", error="changed", goal_verdict={"satisfied": False}, total_tokens=5)
row = await store.get("run")
assert row.get("error") is None
assert row["goal_verdict"] == verdict
assert row["total_tokens"] == 5
@pytest.mark.anyio
async def test_terminal_cancel_row_does_not_report_active_cancellation_win(store):
await put(store)
assert await store.request_cancel("run", action="interrupt") == "interrupt"
assert await store.update_status("run", "interrupted", owner_worker_id="owner")
result = await store.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner")
assert not result.finalized
assert result.cancel_action is None
@pytest.mark.anyio
async def test_owned_status_success_cannot_bypass_durable_cancel(store):
await put(store)
assert await store.request_cancel("run", action="interrupt") == "interrupt"
assert await store.update_status("run", "success", owner_worker_id="owner") is False
assert (await store.get("run"))["status"] == "running"