Totoro cd0e74edaf
fix(scheduler): reconcile stuck once tasks from committed run outcome (#5035)
* fix(scheduler): reconcile stuck once tasks from committed run outcome

Restart recovery (cancel_stuck_once_tasks and the multi-instance
reconcile_stuck_once_tasks) blindly flipped every stuck once-task to
'cancelled'. When handle_run_completion crashed between its two
transactions, a once-task whose run had already committed 'success'
was permanently reported as cancelled.

Both reconciliation paths now read the latest scheduled_task_runs row
without a status filter and finalize the parent to match:
success -> completed (last_error cleared),
failed -> failed with the run's error,
interrupted -> cancelled with the run's error when present,
skipped -> cancelled (no work performed).
Active occurrences (queued/launching/running) are left untouched — a
concurrent completion or a later recovery pass will finalize them once
the run reaches a terminal state.
Tasks without a terminal run row keep the previous generic cancellation.

Review follow-ups (willem-bd / Huixin615):
- Extract _finalise_once_task_from_run() so both recovery paths share one
  outcome mapping (no more drift between single- and multi-instance paths).
  Returns bool (True = finalised, False = active/no-op) for explicit
  counter management at call sites.
- Fix a no-op (`run_row.error or None` -> `run_row.error`) in the skipped
  branch.
- Drop the unused `status` parameter from the test task helpers.
- Use TERMINAL_RUN_STATUSES / ACTIVE_RUN_STATUSES constants (local copies
  to avoid circular import; kept in sync with scheduled_task_runs.sql).
- [P1] Read the latest run AFTER acquiring the parent task row lock, not from
  a pre-lock batch snapshot. The latest-run lookup now runs per task under
  the lock with populate_existing so a concurrently committed status is read
  back fresh.
- [P2] Race tests now use monkeypatch to actually enter the race window:
  _intercepted_fetch commits success in a separate session at the moment the
  per-task fetch fires, so a reverted pre-lock batch implementation fails the
  test, while the current post-lock implementation passes.
- [P1] Do not finalize parent for active occurrences. A non-terminal
  scheduled occurrence means the run is still in progress — the parent must
  be left untouched until the completion path or a later recovery pass
  establishes a terminal outcome.
- [P2] Add cancel_stuck_once_tasks to the single-instance poll loop so
  stuck once-tasks are not left permanently "running" when the startup sweep
  fails (mirrors multi-instance _reconcile_active_state behavior).
- Fix stale docstrings in cancel_stuck_once_tasks and _fetch_latest_run.

Adds regression tests for multiple historical runs (older success +
newer skipped/active) on both paths, monkeypatch-based race tests that
prove a concurrent completion committing success is reflected as
completed, and active-run tests that verify the parent is left
unchanged. Documents the behavior in AGENTS.md.

Fixes #5034

* fix(scheduler): address review comments on completion-consistency fix

- _fetch_latest_run: drop arbitrary id DESC tie-break; order by
  scheduled_for DESC (deterministic recency on schedule position)
- _finalise_once_task_from_run: annotate bool return type
- Centralize TERMINAL/ACTIVE_RUN_STATUSES in scheduled_tasks/model.py;
  stop duplicating them in scheduled_tasks/sql.py and
  scheduled_task_runs/sql.py (removes stale circular-import workaround)
- cancel_stuck_once_tasks: run unconditionally in single-instance poll
  loop (remove try/except swallow)
- tests: pin created_at/scheduled_for in _create_run so recency ordering
  is actually exercised; correct docstrings that described the
  active-occurrence branch as 'generic cancel' instead of 'left unchanged'

* fix(scheduler): correct finalizer return annotation

* fix: order scheduled task runs by creation time

* fix(scheduler): stabilize latest run reconciliation ordering

* fix(scheduler): order latest runs by creation time

* test: update trace scheduler stub

* fix(scheduler): clarify reconciliation diagnostics

Signed-off-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com>

* fix(scheduler): fail closed on startup recovery

Keep single-instance parent reconciliation at startup so it cannot race manual admission. Propagate recovery failures through the Gateway lifespan before channel startup, preventing a half-started scheduler.

Tests cover both recovery failure stages and a queued occurrence that survives startup before the ordinary poll drain launches it.

Signed-off-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com>

* fix(scheduler): order occurrences and fence stale parent writes

Allocate per-task occurrence sequences under the parent lock and guard parent projection across launch, recovery, completion, and queue failure paths. Track launch accounting separately so stale occurrences are counted once without replacing newer results. Commit completion and accounting atomically, preserve legacy history, and cover migrations and reordered execution on SQLite and PostgreSQL.

* fix(scheduler): tighten completion projection and launch fencing diagnostics

Share the once-task outcome mapping between completion and both recovery paths, validate the terminal status before opening the completion transaction, leave cron parent status untouched on completion, log the fenced launch update when an occurrence does not belong to the launched run, and drop the README capability line.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* fix(scheduler): compare caller time only against unsequenced occurrences

Among sequenced rows the parent-locked occurrence_seq is the only recency key. An unsequenced row can only be legacy history or an admission by a pre-upgrade Gateway writer, so recovery prefers it over the sequence winner only when its caller timestamp is later, which is the previous ordering for that pair. A rolling upgrade therefore degrades to the pre-sequence behaviour instead of ranking every pre-upgrade admission below every sequenced one. Document that boundary instead of requiring every Gateway writer to stop before the upgrade.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* fix(scheduler): gate once-task recovery on the same projection rule

Recovery now finalises a once-task parent only from the occurrence that can_project() accepts: the highest sequenced occurrence whenever one exists, or the timestamp-latest row for a task whose history is entirely unsequenced. An unsequenced row admitted by a pre-upgrade writer can no longer cancel a parent whose sequenced occurrence is still live, nor stall finalisation of a parent whose sequenced occurrence already completed. Document that pre-upgrade instances project their own admissions during a rolling upgrade.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* fix(scheduler): defer once-task recovery while any occurrence is live

uq_scheduled_task_run_active allows one non-terminal occurrence per task, so a live row is the task's newest admission whatever its caller clock and whether it carries a sequence. Both once-task recovery paths now probe for any active occurrence after the fresh latest-run read and leave the parent untouched while one exists; cancel_stuck_once_tasks also locks the parent row so admission cannot insert a queued occurrence between that probe and the commit. Once no occurrence is live, the sequence winner decides and a terminalised unsequenced row never overrides it.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* test(persistence): follow the local head past canonical 0019

Main's forward-revision tests assumed 0019_thread_incarnations was the local chain head. With 0022_scheduled_occurrence_seq chained after it, seed the canonical-0019 shape explicitly, assert the real head where a database is upgraded, derive the 0020 rollback binary's revision set from the ancestors of its head, and step the PostgreSQL restart scenario back to canonical 0019 before the rollback binary restarts.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* docs(migrations): describe the chain through 0022_scheduled_occurrence_seq

The rolling-forward section still ended the local chain at canonical 0019; it now names 0022_scheduled_occurrence_seq as the head and lists it among the revisions the 0020 rollback-floor binary does not know.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* test(scheduler): accept CI's sync Postgres URL in occurrence fixtures

CI hands over TEST_POSTGRES_URI as postgresql://...?sslmode=disable. The occurrence, ordering and 0022 migration fixtures built async engines from it directly, so SQLAlchemy chose psycopg2, which is not installed. Normalize the scheme to postgresql+asyncpg and drop libpq-only query keys, matching the existing 0019 migration tests.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* docs(scheduler): keep the backend AGENTS.md chain within its budget

The middlewares guidance chain was already above the hard limit on main, so any added byte in backend/AGENTS.md fails the agent guidance check. Leave backend/AGENTS.md identical to main and record the recovery projection rule in the 0022 migration entry, which already describes the occurrence fields.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

---------

Signed-off-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com>
Co-authored-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-12 07:35:07 +08:00

843 lines
37 KiB
Python

from __future__ import annotations
from datetime import UTC, datetime, timedelta
from typing import Any
from sqlalchemy import and_, exists, func, or_, select, text, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from sqlalchemy.orm import aliased
from deerflow.persistence.run import RunRepository
from deerflow.persistence.run.model import RunRow
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_task_runs.projection import account_launch, can_project
from deerflow.persistence.scheduled_tasks.model import ACTIVE_RUN_STATUSES, TERMINAL_RUN_STATUSES, ScheduledTaskRow
from deerflow.scheduler.schedules import next_run_at as compute_next_run_at
from deerflow.utils.time import coerce_iso
EXECUTING_RUN_STATUSES: tuple[str, ...] = ("launching", "running")
_SCHEDULER_BUDGET_LOCK_KEY = 4694001
def _lease_is_alive(lease_expires_at: datetime | None, *, now: datetime, grace_seconds: int) -> bool:
if lease_expires_at is None:
return False
if lease_expires_at.tzinfo is None:
lease_expires_at = lease_expires_at.replace(tzinfo=UTC)
return lease_expires_at >= now - timedelta(seconds=grace_seconds)
class ActiveScheduledRunConflict(Exception):
"""A concurrent dispatch already holds the task's active-occurrence slot.
Coordinated admission serializes on the parent task before inserting the
queue row. The partial unique index remains the database backstop for
direct repository callers and legacy interleavings. Both paths surface the
same domain exception without coupling the service to SQLAlchemy errors.
"""
def __init__(self, task_id: str) -> None:
self.task_id = task_id
super().__init__(f"scheduled task {task_id!r} already has an active run")
class ScheduledTaskAdmissionRejected(Exception):
"""The parent task changed or disappeared before queue admission."""
def __init__(self, task_id: str, *, reason: str) -> None:
self.task_id = task_id
self.reason = reason
super().__init__(f"scheduled task {task_id!r} admission rejected: {reason}")
class ScheduledTaskRunRepository:
def __init__(
self,
session_factory: async_sessionmaker[AsyncSession],
*,
run_repository: RunRepository | None = None,
) -> None:
self._sf = session_factory
self._run_repository = run_repository or RunRepository(session_factory)
@staticmethod
def _row_to_dict(row: ScheduledTaskRunRow) -> dict[str, Any]:
data = row.to_dict(exclude={"occurrence_seq", "launch_accounted"})
for key in (
"scheduled_for",
"lease_expires_at",
"started_at",
"finished_at",
"created_at",
):
if data.get(key) is not None:
data[key] = coerce_iso(data[key])
return data
@staticmethod
async def _lock_task(session: AsyncSession, task_id: str) -> ScheduledTaskRow | None:
# SQLite ignores SELECT ... FOR UPDATE. Touch the parent first so its
# single-writer lock provides the same serialization point used by
# Postgres row locking for admission, mutation, pause, and delete.
if session.get_bind().dialect.name == "sqlite":
await session.execute(update(ScheduledTaskRow).where(ScheduledTaskRow.id == task_id).values(updated_at=ScheduledTaskRow.updated_at))
return await session.get(ScheduledTaskRow, task_id, with_for_update=True)
@staticmethod
def _associate_scheduled_run(
row: ScheduledTaskRunRow,
candidate: RunRow,
) -> None:
"""Fill launch bookkeeping that the durable run already proves."""
row.run_id = candidate.run_id
if row.started_at is None:
row.started_at = candidate.created_at
@staticmethod
def _associate_task_with_run(
task: ScheduledTaskRow | None,
row: ScheduledTaskRunRow,
candidate: RunRow,
) -> None:
"""Repair the parent update if launch committed before bookkeeping."""
if task is None:
return
account_launch(task, row, candidate.run_id)
if task.last_run_id == candidate.run_id or not can_project(task, row):
return
launched_at = candidate.created_at
if launched_at.tzinfo is None:
launched_at = launched_at.replace(tzinfo=UTC)
task.last_run_at = launched_at
task.last_run_id = candidate.run_id
task.last_thread_id = row.thread_id
task.next_run_at = compute_next_run_at(
task.schedule_type,
task.schedule_spec,
task.timezone,
now=launched_at,
)
task.lease_owner = None
task.lease_expires_at = None
task.updated_at = datetime.now(UTC)
if task.schedule_type == "once":
if candidate.status == "success":
task.status = "completed"
task.last_error = None
elif candidate.status in {"error", "timeout"}:
task.status = "failed"
task.last_error = candidate.error
elif candidate.status == "interrupted":
task.status = "cancelled"
task.last_error = candidate.error
else:
task.status = "running"
task.last_error = None
elif not (row.trigger == "manual" and task.status == "paused"):
task.status = "enabled"
task.last_error = candidate.error if candidate.status in {"error", "timeout", "interrupted"} else None
async def create(
self,
*,
run_record_id: str,
task_id: str,
thread_id: str,
scheduled_for: datetime,
trigger: str,
status: str,
coordinate_with_task: bool = False,
expected_task_user_id: str | None = None,
expected_task_status: str | None = None,
expected_task_updated_at: datetime | str | None = None,
expected_task_lease_owner: str | None = None,
release_task_lease_status: str | None = None,
) -> dict[str, Any]:
row = ScheduledTaskRunRow(
id=run_record_id,
task_id=task_id,
thread_id=thread_id,
scheduled_for=scheduled_for,
trigger=trigger,
status=status,
launch_accounted=False,
created_at=datetime.now(UTC),
)
async with self._sf() as session:
# Every occurrence with a parent participates in the same DB
# ordering, including terminal records and uncoordinated callers.
task = await self._lock_task(session, task_id)
if coordinate_with_task:
if task is None or (expected_task_user_id is not None and task.user_id != expected_task_user_id):
await session.rollback()
raise ScheduledTaskAdmissionRejected(task_id, reason="not_found")
if expected_task_lease_owner is not None:
if task.lease_owner != expected_task_lease_owner:
await session.rollback()
raise ScheduledTaskAdmissionRejected(task_id, reason="stale")
else:
if expected_task_status is not None and task.status != expected_task_status:
await session.rollback()
raise ScheduledTaskAdmissionRejected(task_id, reason="stale")
if expected_task_updated_at is not None and coerce_iso(task.updated_at) != coerce_iso(expected_task_updated_at):
await session.rollback()
raise ScheduledTaskAdmissionRejected(task_id, reason="stale")
active_status = await session.scalar(
select(ScheduledTaskRunRow.status)
.where(
ScheduledTaskRunRow.task_id == task_id,
ScheduledTaskRunRow.status.in_(ACTIVE_RUN_STATUSES),
)
.limit(1)
)
if active_status is not None:
await session.rollback()
raise ActiveScheduledRunConflict(task_id)
if task is not None:
row.occurrence_seq = await session.scalar(
update(ScheduledTaskRow)
.where(ScheduledTaskRow.id == task_id)
.values(
last_occurrence_seq=ScheduledTaskRow.last_occurrence_seq + 1,
# Sequence allocation is not a user-visible task edit.
updated_at=ScheduledTaskRow.updated_at,
)
.returning(ScheduledTaskRow.last_occurrence_seq)
)
session.add(row)
if coordinate_with_task and task is not None and release_task_lease_status is not None:
task.status = release_task_lease_status
task.lease_owner = None
task.lease_expires_at = None
task.updated_at = datetime.now(UTC)
try:
await session.commit()
except IntegrityError:
await session.rollback()
# A primary-key/sequence conflict is not necessarily an active
# slot conflict. Preserve the database error unless an active
# occurrence actually exists after rollback.
if status in ACTIVE_RUN_STATUSES and await session.scalar(select(ScheduledTaskRunRow.id).where(ScheduledTaskRunRow.task_id == task_id, ScheduledTaskRunRow.status.in_(ACTIVE_RUN_STATUSES)).limit(1)):
raise ActiveScheduledRunConflict(task_id) from None
raise
await session.refresh(row)
return self._row_to_dict(row)
async def list_by_task(self, task_id: str, *, limit: int = 50, offset: int = 0) -> list[dict[str, Any]]:
stmt = (
select(ScheduledTaskRunRow)
.where(ScheduledTaskRunRow.task_id == task_id)
.order_by(
ScheduledTaskRunRow.created_at.desc(),
ScheduledTaskRunRow.id.desc(),
)
.limit(limit)
.offset(offset)
)
async with self._sf() as session:
result = await session.execute(stmt)
return [self._row_to_dict(row) for row in result.scalars()]
async def count_active_runs(self) -> int:
"""Count launch claims and live runs; waiting rows do not consume slots."""
stmt = select(func.count()).select_from(ScheduledTaskRunRow).where(ScheduledTaskRunRow.status.in_(EXECUTING_RUN_STATUSES))
async with self._sf() as session:
result = await session.execute(stmt)
return int(result.scalar() or 0)
async def list_queued_runs(self, *, limit: int) -> list[dict[str, Any]]:
older = aliased(ScheduledTaskRunRow)
older_same_thread = exists(
select(older.id).where(
older.thread_id == ScheduledTaskRunRow.thread_id,
older.status.in_(ACTIVE_RUN_STATUSES),
or_(
older.created_at < ScheduledTaskRunRow.created_at,
and_(
older.created_at == ScheduledTaskRunRow.created_at,
older.id < ScheduledTaskRunRow.id,
),
),
)
)
stmt = (
select(ScheduledTaskRunRow)
.where(
ScheduledTaskRunRow.status == "queued",
~older_same_thread,
)
# Prefer rows that have had fewer launch attempts. A permanently
# busy thread therefore cannot monopolize the bounded drain batch,
# while created_at/id preserve FIFO order among equal attempts.
.order_by(
ScheduledTaskRunRow.attempt_count.asc(),
ScheduledTaskRunRow.created_at.asc(),
ScheduledTaskRunRow.id.asc(),
)
.limit(limit)
)
async with self._sf() as session:
result = await session.execute(stmt)
return [self._row_to_dict(row) for row in result.scalars()]
async def get_active_run(self, task_id: str) -> dict[str, Any] | None:
stmt = (
select(ScheduledTaskRunRow)
.where(
ScheduledTaskRunRow.task_id == task_id,
ScheduledTaskRunRow.status.in_(ACTIVE_RUN_STATUSES),
)
.order_by(ScheduledTaskRunRow.created_at.asc(), ScheduledTaskRunRow.id.asc())
.limit(1)
)
async with self._sf() as session:
row = (await session.execute(stmt)).scalars().first()
return self._row_to_dict(row) if row is not None else None
async def claim_queued_run(
self,
run_record_id: str,
*,
lease_owner: str,
now: datetime,
lease_seconds: int,
global_max_concurrent_runs: int,
) -> dict[str, Any] | None:
"""Atomically move one waiting row into the lease-fenced launch phase."""
async with self._sf() as session:
if session.get_bind().dialect.name == "postgresql":
await session.execute(
text("SELECT pg_advisory_xact_lock(:lock_key)"),
{"lock_key": _SCHEDULER_BUDGET_LOCK_KEY},
)
executing = await session.scalar(select(func.count()).select_from(ScheduledTaskRunRow).where(ScheduledTaskRunRow.status.in_(EXECUTING_RUN_STATUSES)))
if int(executing or 0) >= global_max_concurrent_runs:
await session.rollback()
return None
older = aliased(ScheduledTaskRunRow)
older_same_thread = exists(
select(older.id).where(
older.thread_id == ScheduledTaskRunRow.thread_id,
older.status.in_(ACTIVE_RUN_STATUSES),
or_(
older.created_at < ScheduledTaskRunRow.created_at,
and_(
older.created_at == ScheduledTaskRunRow.created_at,
older.id < ScheduledTaskRunRow.id,
),
),
)
)
result = await session.execute(
update(ScheduledTaskRunRow)
.where(
ScheduledTaskRunRow.id == run_record_id,
ScheduledTaskRunRow.status == "queued",
~older_same_thread,
)
.values(
status="launching",
lease_owner=lease_owner,
lease_expires_at=now + timedelta(seconds=lease_seconds),
attempt_count=ScheduledTaskRunRow.attempt_count + 1,
)
)
if result.rowcount != 1:
await session.rollback()
return None
await session.commit()
row = await session.get(ScheduledTaskRunRow, run_record_id)
return self._row_to_dict(row) if row is not None else None
async def requeue_claimed_run(
self,
run_record_id: str,
*,
lease_owner: str,
error: str | None = None,
) -> bool:
async with self._sf() as session:
result = await session.execute(
update(ScheduledTaskRunRow)
.where(
ScheduledTaskRunRow.id == run_record_id,
ScheduledTaskRunRow.status == "launching",
ScheduledTaskRunRow.lease_owner == lease_owner,
)
.values(
status="queued",
lease_owner=None,
lease_expires_at=None,
error=error,
)
)
await session.commit()
return result.rowcount == 1
async def expire_queued_runs(
self,
*,
created_before: datetime,
error: str,
now: datetime,
) -> list[dict[str, Any]]:
"""Expire waiting rows and update each parent in one transaction.
The queued child keeps the task unclaimable until the parent lock is
held. Releasing the child slot and advancing the parent in the same
transaction prevents a peer scheduler from taking a stale due-task
lease between those two writes.
"""
async with self._sf() as session:
candidate_keys = list(
(
await session.execute(
select(
ScheduledTaskRunRow.id,
ScheduledTaskRunRow.task_id,
)
.where(
ScheduledTaskRunRow.status == "queued",
ScheduledTaskRunRow.created_at <= created_before,
)
.order_by(
ScheduledTaskRunRow.task_id.asc(),
ScheduledTaskRunRow.id.asc(),
)
)
).all()
)
expired: list[dict[str, Any]] = []
for row_id, task_id in candidate_keys:
async with self._sf() as session:
task = await self._lock_task(session, task_id)
row = await session.get(ScheduledTaskRunRow, row_id, with_for_update=True)
if row is None or row.status != "queued":
await session.rollback()
continue
row.status = "failed"
row.error = error
row.finished_at = now
row.lease_owner = None
row.lease_expires_at = None
if task is not None and can_project(task, row):
if task.status == "paused":
# A concurrent/later pause owns the parent presentation.
task.lease_owner = None
task.lease_expires_at = None
else:
if row.trigger == "manual":
next_at = task.next_run_at
task_status = task.status or "enabled"
else:
next_at = compute_next_run_at(
task.schedule_type,
task.schedule_spec,
task.timezone,
now=now,
)
task_status = "failed" if task.schedule_type == "once" else "enabled"
task.status = task_status
task.next_run_at = next_at
task.last_thread_id = row.thread_id
task.last_error = error
task.lease_owner = None
task.lease_expires_at = None
task.updated_at = datetime.now(UTC)
await session.commit()
expired.append(self._row_to_dict(row))
return expired
async def fail_launching_run(
self,
run_record_id: str,
*,
task_id: str,
lease_owner: str,
error: str,
now: datetime,
) -> bool:
"""Fail a claimed launch and update its parent without a release gap."""
async with self._sf() as session:
task = await self._lock_task(session, task_id)
row = await session.get(ScheduledTaskRunRow, run_record_id, with_for_update=True)
if row is None or row.task_id != task_id or row.status != "launching" or row.lease_owner != lease_owner:
await session.rollback()
return False
row.status = "failed"
row.error = error
row.started_at = now
row.finished_at = now
row.lease_owner = None
row.lease_expires_at = None
if task is not None and can_project(task, row):
if row.trigger == "manual":
task_status = task.status or "enabled"
next_at = task.next_run_at
else:
task_status = "failed" if task.schedule_type == "once" else "enabled"
next_at = compute_next_run_at(
task.schedule_type,
task.schedule_spec,
task.timezone,
now=now,
)
task.status = task_status
task.next_run_at = next_at
task.last_run_at = now
task.last_run_id = None
task.last_thread_id = row.thread_id
task.last_error = error
task.lease_owner = None
task.lease_expires_at = None
task.updated_at = datetime.now(UTC)
await session.commit()
return True
async def reconcile_launched_run(
self,
run_record_id: str,
*,
task_id: str,
run_id: str,
started_at: datetime,
) -> bool:
"""Associate a run that committed after its short claim was recovered.
The Gateway launch path is idempotent per scheduled occurrence, so a
peer that already reclaimed this row resolves to the same durable run.
Parent-first locking restores the active slot before later parent
bookkeeping can make the task due again.
"""
async with self._sf() as session:
await self._lock_task(session, task_id)
row = await session.get(ScheduledTaskRunRow, run_record_id, with_for_update=True)
if row is None or row.task_id != task_id:
await session.rollback()
return False
if row.status in TERMINAL_RUN_STATUSES:
if row.run_id is None:
# Timeout/pause can terminalize a claim that recovery
# briefly returned to ``queued`` while Gateway admission
# was still completing. A returned durable run is the
# stronger fact; completion-owned terminal rows already
# carry this exact run_id and stay terminal below.
row.status = "running"
row.run_id = run_id
row.error = None
row.finished_at = None
row.started_at = row.started_at or started_at
elif row.run_id != run_id:
await session.rollback()
return False
elif row.started_at is None:
row.started_at = started_at
else:
row.status = "running"
row.run_id = run_id
row.error = None
row.started_at = row.started_at or started_at
row.lease_owner = None
row.lease_expires_at = None
await session.commit()
return True
async def recover_expired_launch_claims(self, *, error: str, now: datetime) -> int:
"""Recover single-instance claims that outlived their short lease."""
stmt = (
select(
ScheduledTaskRunRow.id,
ScheduledTaskRunRow.task_id,
)
.where(
ScheduledTaskRunRow.status == "launching",
or_(
ScheduledTaskRunRow.lease_expires_at.is_(None),
ScheduledTaskRunRow.lease_expires_at < now,
),
)
.order_by(
ScheduledTaskRunRow.task_id.asc(),
ScheduledTaskRunRow.id.asc(),
)
)
async with self._sf() as session:
row_keys = list((await session.execute(stmt)).all())
recovered = 0
for row_id, task_id in row_keys:
task = await self._lock_task(session, task_id)
row = await session.get(ScheduledTaskRunRow, row_id, with_for_update=True)
if row is None or row.status != "launching":
continue
candidate = await self._find_underlying_run(session, row, task)
row.lease_owner = None
row.lease_expires_at = None
if candidate is None:
row.status = "queued"
else:
self._associate_scheduled_run(row, candidate)
self._associate_task_with_run(task, row, candidate)
if candidate.status in {"pending", "running"}:
row.status = "running"
row.error = None
elif candidate.status == "success":
row.status = "success"
row.error = None
row.finished_at = now
elif candidate.status in {"error", "timeout"}:
row.status = "failed"
row.error = candidate.error
row.finished_at = now
else:
row.status = "interrupted"
row.error = candidate.error or error
row.finished_at = now
recovered += 1
await session.commit()
return recovered
async def update_status(
self,
run_record_id: str,
*,
status: str,
run_id: str | None = None,
error: str | None = None,
started_at: datetime | None = None,
finished_at: datetime | None = None,
protect_terminal: bool = False,
expected_lease_owner: str | None = None,
) -> bool:
async with self._sf() as session:
row = await session.get(ScheduledTaskRunRow, run_record_id)
if row is None:
return False
if protect_terminal and row.status in TERMINAL_RUN_STATUSES:
# The launch-path "running" write lost the race against the
# completion hook; keep the terminal status/error and only
# backfill bookkeeping the completion write could not know.
# Completion clears the short launch lease, so allow that
# backfill after an owner mismatch only when the terminal row
# already identifies the exact same durable run. A stale
# launcher for another run remains fenced.
same_run = run_id is not None and row.run_id == run_id
if expected_lease_owner is not None and row.lease_owner != expected_lease_owner and not same_run:
await session.rollback()
return False
if row.run_id is None and run_id is not None:
row.run_id = run_id
if row.started_at is None and started_at is not None:
row.started_at = started_at
await session.commit()
return True
if expected_lease_owner is not None and row.lease_owner != expected_lease_owner:
await session.rollback()
return False
row.status = status
row.run_id = run_id
row.error = error
if status != "launching":
row.lease_owner = None
row.lease_expires_at = None
if started_at is not None:
row.started_at = started_at
if finished_at is not None:
row.finished_at = finished_at
await session.commit()
return True
async def has_active_runs(self, task_id: str) -> bool:
stmt = (
select(ScheduledTaskRunRow.id)
.where(
ScheduledTaskRunRow.task_id == task_id,
ScheduledTaskRunRow.status.in_(ACTIVE_RUN_STATUSES),
)
.limit(1)
)
async with self._sf() as session:
result = await session.execute(stmt)
return result.scalars().first() is not None
async def mark_stale_active_runs(self, *, error: str) -> int:
"""Recover single-instance launch claims and fail orphaned live runs.
Waiting rows are durable queue entries and survive restart. A
``launching`` row without a committed live run is safe to retry; a
``running`` row belonged to the dead in-process runtime.
"""
stmt = (
select(
ScheduledTaskRunRow.id,
ScheduledTaskRunRow.task_id,
)
.where(ScheduledTaskRunRow.status.in_(EXECUTING_RUN_STATUSES))
.order_by(
ScheduledTaskRunRow.task_id.asc(),
ScheduledTaskRunRow.id.asc(),
)
)
now = datetime.now(UTC)
async with self._sf() as session:
row_keys = list((await session.execute(stmt)).all())
recovered = 0
for row_id, task_id in row_keys:
task = await self._lock_task(session, task_id)
row = await session.get(ScheduledTaskRunRow, row_id, with_for_update=True)
if row is None or row.status not in EXECUTING_RUN_STATUSES:
continue
row.lease_owner = None
row.lease_expires_at = None
candidate = await self._find_underlying_run(session, row, task)
if row.status == "launching" and candidate is None:
row.status = "queued"
else:
if candidate is not None:
self._associate_scheduled_run(row, candidate)
self._associate_task_with_run(task, row, candidate)
if candidate is not None and candidate.status == "success":
row.status = "success"
row.error = None
elif candidate is not None and candidate.status in {"error", "timeout"}:
row.status = "failed"
row.error = candidate.error
else:
row.status = "interrupted"
row.error = error
row.finished_at = now
recovered += 1
await session.commit()
return recovered
async def reconcile_active_runs(
self,
*,
error: str,
now: datetime,
lease_grace_seconds: int = 10,
) -> int:
"""Reconcile only rows whose underlying owner is no longer live.
``RunManager`` owns the durable run lease. A scheduled row with a live
underlying run, or a queued row whose parent task still has a dispatch
lease, belongs to another process and must survive this startup.
"""
async with self._sf() as session:
result = await session.execute(
select(
ScheduledTaskRunRow.id,
ScheduledTaskRunRow.task_id,
)
.where(ScheduledTaskRunRow.status.in_(EXECUTING_RUN_STATUSES))
.order_by(
ScheduledTaskRunRow.task_id.asc(),
ScheduledTaskRunRow.id.asc(),
)
)
row_keys = list(result.all())
stale = 0
associations: list[tuple[ScheduledTaskRow | None, ScheduledTaskRunRow, RunRow]] = []
for row_id, task_id in row_keys:
# Keep the same task -> scheduled-run lock order used by
# pause/delete. Reversing these two locks lets a user action
# and a peer reconciliation deadlock each other on Postgres.
# Multi-instance reconciliation is Postgres-only. Keep this a
# row lock without SQLite's writer-lock emulation because the
# SQLite regression path performs durable-run takeover in a
# nested short transaction below.
task = await session.get(ScheduledTaskRow, task_id, with_for_update=True)
row = await session.get(ScheduledTaskRunRow, row_id, with_for_update=True)
if row is None or row.status not in EXECUTING_RUN_STATUSES:
continue
candidate = await self._find_underlying_run(session, row, task)
if candidate is not None:
self._associate_scheduled_run(row, candidate)
# Defer parent writes until all run takeovers have finished.
# Flushing a parent mutation before claim_for_takeover()
# would hold SQLite's writer lock across the nested short
# transaction used by that durable-run CAS.
associations.append((task, row, candidate))
if candidate is not None and candidate.status not in {"pending", "running"}:
row.lease_owner = None
row.lease_expires_at = None
if candidate.status == "success":
row.status = "success"
row.error = None
elif candidate.status in {"error", "timeout"}:
row.status = "failed"
row.error = candidate.error
else:
row.status = "interrupted"
row.error = candidate.error or error
row.finished_at = now
stale += 1
continue
if candidate is not None and candidate.status in {"pending", "running"}:
if _lease_is_alive(candidate.lease_expires_at, now=now, grace_seconds=lease_grace_seconds):
# A peer can observe the committed durable run before
# the launcher writes its scheduled-run bookkeeping.
# Once reconciliation releases the short launch claim,
# that owner-fenced write must be allowed to fail
# without leaving incomplete or stale history behind.
row.error = None
if row.status == "launching":
row.status = "running"
row.lease_owner = None
row.lease_expires_at = None
continue
# Run takeover commits in its own short transaction. If this
# outer commit fails, the next poll finishes scheduled-row
# bookkeeping while the run remains safely terminal.
claimed = await self._run_repository.claim_for_takeover(
candidate.run_id,
grace_seconds=lease_grace_seconds,
error=error,
stop_reason="scheduled_task_orphan_recovered",
)
if not claimed:
refreshed = await self._run_repository.get(candidate.run_id, user_id=None)
if refreshed is not None and refreshed.get("status") in {"pending", "running"}:
continue
if row.status == "launching" and row.run_id is None:
if _lease_is_alive(row.lease_expires_at, now=now, grace_seconds=0):
continue
row.status = "queued"
row.lease_owner = None
row.lease_expires_at = None
stale += 1
continue
row.status = "interrupted"
row.error = error
row.finished_at = now
row.lease_owner = None
row.lease_expires_at = None
stale += 1
for task, row, candidate in associations:
self._associate_task_with_run(task, row, candidate)
await session.commit()
return stale
@staticmethod
async def _find_underlying_run(session: AsyncSession, row: ScheduledTaskRunRow, task: ScheduledTaskRow | None) -> RunRow | None:
run_ids = [candidate for candidate in (row.run_id, task.last_run_id if task is not None else None) if candidate]
for run_id in dict.fromkeys(run_ids):
candidate = await session.get(RunRow, run_id)
if candidate is None:
continue
linked_task_run_id = (candidate.metadata_json or {}).get("scheduled_task_run_id")
# A stale parent ``last_run_id`` may point at a previous occurrence.
# Let the current scheduled-run metadata lookup recover the live row.
if linked_task_run_id is None or linked_task_run_id == row.id:
return candidate
result = await session.execute(select(RunRow).where(RunRow.metadata_json["scheduled_task_run_id"].as_string() == row.id).order_by(RunRow.created_at.desc()).limit(1))
return result.scalars().first()