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>
This commit is contained in:
Totoro 2026-09-12 07:35:07 +08:00 committed by GitHub
parent 9f17bbeec7
commit cd0e74edaf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
21 changed files with 2727 additions and 175 deletions

View File

@ -1565,6 +1565,7 @@ The background scheduler is single-instance by default. For a multi-pod deployme
### Upgrade Notes ### Upgrade Notes
- Occurrence ordering applies to rows admitted by upgraded Gateway instances, which project only sequenced occurrences onto the parent task and defer recovery while any occurrence is still live, whichever instance admitted it; a task whose history is entirely unsequenced keeps the previous timestamp ordering until its first sequenced admission. During a rolling upgrade, rows admitted by pre-upgrade instances are projected by those instances themselves, as before the upgrade, and the ordering guarantees hold once every Gateway writer runs the upgraded version. Existing history is not backfilled; the upgrade does not reconstruct past order or repair historical counts.
- Before upgrading a deployment with `GATEWAY_WORKERS > 1` and `scheduler.enabled: true`, either keep the scheduler on exactly one Gateway worker or configure `scheduler.multi_instance: true` with shared Postgres, `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`. The upgraded Gateway rejects the unsafe combination at startup instead of starting silently. - Before upgrading a deployment with `GATEWAY_WORKERS > 1` and `scheduler.enabled: true`, either keep the scheduler on exactly one Gateway worker or configure `scheduler.multi_instance: true` with shared Postgres, `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`. The upgraded Gateway rejects the unsafe combination at startup instead of starting silently.
- In multi-instance mode, `scheduler.max_concurrent_runs` is a cluster-wide execution cap, not a per-Pod cap. It includes `launching` and `running` scheduled occurrences, so capacity does not multiply with the number of replicas; durable waiting rows remain outside the cap. - In multi-instance mode, `scheduler.max_concurrent_runs` is a cluster-wide execution cap, not a per-Pod cap. It includes `launching` and `running` scheduled occurrences, so capacity does not multiply with the number of replicas; durable waiting rows remain outside the cap.
- `scheduler.multi_instance` and the related scheduler, ownership, and run-event settings are startup-only. Apply changes with a coordinated restart of all Gateway Pods; changing the ConfigMap alone does not activate multi-instance recovery. - `scheduler.multi_instance` and the related scheduler, ownership, and run-event settings are startup-only. Apply changes with a coordinated restart of all Gateway Pods; changing the ConfigMap alone does not activate multi-instance recovery.

View File

@ -306,27 +306,6 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
# Must run AFTER langgraph_runtime so app.state.store is available for thread migration # Must run AFTER langgraph_runtime so app.state.store is available for thread migration
await _ensure_admin_user(app) await _ensure_admin_user(app)
# Start IM channel service if any channels are configured
try:
from app.channels.service import start_channel_service
# Closure over `app` (mirrors ScheduledTaskService's `launch_run`
# below) rather than resolving `app.state.stream_bridge` here
# directly: `stream_bridge` is a STARTUP_ONLY_FIELDS singleton set
# once, above, by `langgraph_runtime(app, startup_config)`, so
# either shape is safe by construction — the closure is just the
# more defensive/consistent-with-precedent form, and it is what
# ChannelManager's follow-up-drain watcher (issue #4121 Slice 2)
# uses to reach the same StreamBridge every other run consumer
# goes through `get_stream_bridge(request)` for.
channel_service = await start_channel_service(
startup_config,
get_stream_bridge=lambda: getattr(app.state, "stream_bridge", None),
)
logger.info("Channel service started: %s", channel_service.get_status())
except Exception:
logger.exception("No IM channels configured or channel service failed to start")
try: try:
from app.gateway.services import launch_scheduled_thread_run from app.gateway.services import launch_scheduled_thread_run
from app.scheduler import ScheduledTaskService from app.scheduler import ScheduledTaskService
@ -348,6 +327,33 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
await scheduled_task_service.start() await scheduled_task_service.start()
except Exception: except Exception:
logger.exception("Failed to initialize scheduled task service") logger.exception("Failed to initialize scheduled task service")
# If an enabled scheduler rejects start(), keep that rejection as a
# lifespan failure instead of exposing a half-started service.
if startup_config.scheduler.enabled:
raise
# Start IM channel service only after scheduler recovery succeeds, so a
# fail-closed scheduler startup cannot strand channel-owned tasks before
# the lifespan reaches its normal shutdown boundary.
try:
from app.channels.service import start_channel_service
# Closure over `app` (mirrors ScheduledTaskService's `launch_run`
# above) rather than resolving `app.state.stream_bridge` here
# directly: `stream_bridge` is a STARTUP_ONLY_FIELDS singleton set
# once, above, by `langgraph_runtime(app, startup_config)`, so
# either shape is safe by construction — the closure is just the
# more defensive/consistent-with-precedent form, and it is what
# ChannelManager's follow-up-drain watcher (issue #4121 Slice 2)
# uses to reach the same StreamBridge every other run consumer
# goes through `get_stream_bridge(request)` for.
channel_service = await start_channel_service(
startup_config,
get_stream_bridge=lambda: getattr(app.state, "stream_bridge", None),
)
logger.info("Channel service started: %s", channel_service.get_status())
except Exception:
logger.exception("No IM channels configured or channel service failed to start")
from app.gateway.services import launch_mcp_task_notification_run from app.gateway.services import launch_mcp_task_notification_run
from app.mcp_tasks import McpTaskService from app.mcp_tasks import McpTaskService

View File

@ -299,6 +299,7 @@ class ScheduledTaskService:
last_thread_id=launched_thread_id, last_thread_id=launched_thread_id,
last_error=None, last_error=None,
increment_run_count=True, increment_run_count=True,
task_run_id=task_run_id,
# Same race as the run-row write above: a fast-failing run's # Same race as the run-row write above: a fast-failing run's
# completion hook may have already finalized a `once` task. # completion hook may have already finalized a `once` task.
protect_terminal=True, protect_terminal=True,
@ -368,6 +369,7 @@ class ScheduledTaskService:
# The transient itself is logged above. # The transient itself is logged above.
last_error=None, last_error=None,
increment_run_count=True, increment_run_count=True,
task_run_id=task_run_id,
protect_terminal=True, protect_terminal=True,
) )
except Exception: except Exception:
@ -538,31 +540,16 @@ class ScheduledTaskService:
if terminal_status is None: if terminal_status is None:
return return
await self._task_run_repo.update_status( await self._task_repo.complete_run(
task_run_id, task_id,
status=terminal_status, user_id=user_id,
task_run_id=task_run_id,
run_id=record.run_id, run_id=record.run_id,
status=terminal_status,
error=error, error=error,
finished_at=datetime.now(UTC), finished_at=datetime.now(UTC),
) )
task = await self._task_repo.get(task_id, user_id=user_id)
if task is None:
return
updates: dict[str, Any] = {"last_error": error}
if task["schedule_type"] == "once":
# The single occurrence is consumed either way (the run did launch,
# so re-arming risks duplicate side effects), but an interrupt ends
# as "cancelled", not "failed".
if terminal_status == "success":
updates["status"] = "completed"
elif terminal_status == "interrupted":
updates["status"] = "cancelled"
else:
updates["status"] = "failed"
await self._task_repo.update(task_id, user_id=user_id, updates=updates)
async def start(self) -> None: async def start(self) -> None:
if self._task is not None: if self._task is not None:
return return
@ -571,21 +558,28 @@ class ScheduledTaskService:
await self._reconcile_active_state(now=datetime.now(UTC)) await self._reconcile_active_state(now=datetime.now(UTC))
self._skip_next_lease_reconciliation = True self._skip_next_lease_reconciliation = True
else: else:
# This destructive sweep is safe only while Gateway lifespan awaits
# start(): no request or poll admission can create a run owned by
# this process yet. Complete occurrence -> parent recovery before
# returning; moving either pass into run_once() can interrupt live
# work or race manual admission.
try: try:
stale = await self._task_run_repo.mark_stale_active_runs(error=restart_error) stale = await self._task_run_repo.mark_stale_active_runs(error=restart_error)
if stale: if stale:
logger.warning("Marked %d stale scheduled task run(s) as interrupted after restart", stale) logger.warning("Marked %d stale scheduled task run(s) as interrupted after restart", stale)
except Exception: except Exception:
logger.exception("Failed to sweep stale scheduled task runs at startup") logger.exception("Failed to sweep stale scheduled task runs at startup")
raise
try: try:
# The run rows above are only half the story: a launched `once` # The run rows above are only half the story: a launched `once`
# task is parked in "running" until the (now dead) completion hook # task is parked in "running" until the (now dead) completion hook
# would have finalized it, so reconcile the parent rows too. # would have finalized it.
stuck = await self._task_repo.cancel_stuck_once_tasks(error=restart_error) stuck = await self._task_repo.cancel_stuck_once_tasks(error=restart_error)
if stuck: if stuck:
logger.warning("Cancelled %d stuck once task(s) after restart", stuck) logger.warning("Reconciled %d stuck once task(s) after restart", stuck)
except Exception: except Exception:
logger.exception("Failed to reconcile stuck once tasks at startup") logger.exception("Failed to reconcile stuck once tasks at startup")
raise
self._stop.clear() self._stop.clear()
self._task = asyncio.create_task(self._run_loop()) self._task = asyncio.create_task(self._run_loop())
@ -608,7 +602,7 @@ class ScheduledTaskService:
lease_grace_seconds=self._run_lease_grace_seconds, lease_grace_seconds=self._run_lease_grace_seconds,
) )
if stuck: if stuck:
logger.warning("Cancelled %d stuck once task(s) after lease reconciliation", stuck) logger.warning("Reconciled %d stuck once task(s) after lease reconciliation", stuck)
except Exception: except Exception:
logger.exception("Failed to reconcile once tasks with leases") logger.exception("Failed to reconcile once tasks with leases")

View File

@ -22,13 +22,15 @@ The empty-DB path keeps using `create_all` because `Base.metadata` is the only a
**Rolling forward compatibility**: the local chain is **Rolling forward compatibility**: the local chain is
`0018_oauth_identity_pg_partial``0019_projects` `0018_oauth_identity_pg_partial``0019_projects`
`0020_threads_meta_project_id``0021_batch_acceptance` `0020_threads_meta_project_id``0021_batch_acceptance`
`0019_thread_incarnations`. The final revision deliberately retains the exact `0019_thread_incarnations``0022_scheduled_occurrence_seq` (current head).
id audited by the rollback-floor binary; Alembic orders revisions by The incarnation revision deliberately retains the exact id audited by the
`down_revision`, not by the numeric prefix. rollback-floor binary; Alembic orders revisions by `down_revision`, not by the
numeric prefix.
The deployed `0020_threads_meta_project_id` rollback-floor binary knows neither The deployed `0020_threads_meta_project_id` rollback-floor binary knows none of
`0021_batch_acceptance` nor `0019_thread_incarnations`. It treats only the final `0021_batch_acceptance`, `0019_thread_incarnations`, or
incarnation revision as forward-compatible, after reflection confirms every `0022_scheduled_occurrence_seq`. It treats only the incarnation revision as
forward-compatible, after reflection confirms every
table and column in its own ORM schema. The intervening acceptance columns and table and column in its own ORM schema. The intervening acceptance columns and
the incarnation columns are nullable and have no server default, so old the incarnation columns are nullable and have no server default, so old
repositories may omit them. Tests must prove old reads and writes across both repositories may omit them. Tests must prove old reads and writes across both
@ -150,7 +152,8 @@ on installs that never enabled it. The convention is:
- `migrations/versions/0019_projects.py` — creates the `projects` table (id/user_id/name/instructions/presentation/status + timestamps) for the Projects Phase-1 organization feature; chains after `0018_oauth_identity_pg_partial` - `migrations/versions/0019_projects.py` — creates the `projects` table (id/user_id/name/instructions/presentation/status + timestamps) for the Projects Phase-1 organization feature; chains after `0018_oauth_identity_pg_partial`
- `migrations/versions/0020_threads_meta_project_id.py` — adds nullable `threads_meta.project_id` plus `ix_threads_meta_project_id` (no FK by design: project delete clears membership first, and the reserved `deerflow_project_id` metadata key stays in sync); chains after `0019_projects` - `migrations/versions/0020_threads_meta_project_id.py` — adds nullable `threads_meta.project_id` plus `ix_threads_meta_project_id` (no FK by design: project delete clears membership first, and the reserved `deerflow_project_id` metadata key stays in sync); chains after `0019_projects`
- `migrations/versions/0021_batch_acceptance.py` — adds nullable per-item acceptance criteria and verdict JSON columns after `0020_threads_meta_project_id`; legacy rows remain unchecked - `migrations/versions/0021_batch_acceptance.py` — adds nullable per-item acceptance criteria and verdict JSON columns after `0020_threads_meta_project_id`; legacy rows remain unchecked
- `migrations/versions/0019_thread_incarnations.py` — current head; chains after `0021_batch_acceptance` while retaining the exact revision id audited by the rollback-floor binary. Adds nullable `threads_meta.incarnation` / `mcp_tasks.thread_incarnation` columns. New thread rows get a random 32-character incarnation. Memory mutations serialize per thread; an overwrite inherits the existing incarnation, while a delete/recreate gets a new one. SQLite MCP task INSERTs copy the owner-or-shared incarnation with a scalar subquery in the same statement. PostgreSQL task creation holds `FOR SHARE`, which conflicts with both current `FOR UPDATE` mutations and an older writer's plain owner update (`FOR NO KEY UPDATE`). Missing or differently owned threads store NULL, old writers may omit both columns, and current API/task serialization hides them. The migration preflights both tables before DDL and its SQLite downgrade cleans only safe remnants from its own interrupted batch-copy attempt - `migrations/versions/0019_thread_incarnations.py` — chains after `0021_batch_acceptance` while retaining the exact revision id audited by the rollback-floor binary. Adds nullable `threads_meta.incarnation` / `mcp_tasks.thread_incarnation` columns. New thread rows get a random 32-character incarnation. Memory mutations serialize per thread; an overwrite inherits the existing incarnation, while a delete/recreate gets a new one. SQLite MCP task INSERTs copy the owner-or-shared incarnation with a scalar subquery in the same statement. PostgreSQL task creation holds `FOR SHARE`, which conflicts with both current `FOR UPDATE` mutations and an older writer's plain owner update (`FOR NO KEY UPDATE`). Missing or differently owned threads store NULL, old writers may omit both columns, and current API/task serialization hides them. The migration preflights both tables before DDL and its SQLite downgrade cleans only safe remnants from its own interrupted batch-copy attempt
- `migrations/versions/0022_scheduled_occurrence_seq.py` — adds the per-task `last_occurrence_seq` high-water mark, nullable occurrence `occurrence_seq` and `launch_accounted`, and a unique `(task_id, occurrence_seq)` index. New occurrences allocate their sequence under the existing parent lock; launch accounting is recorded atomically with the count so an older recovered occurrence cannot be counted twice. Legacy child columns remain NULL without guessed ordering or accounting backfill. All three fields are internal and omitted from repository responses. Both once-task recovery paths lock the parent, defer while any occurrence row is active (sequenced or not), and otherwise project only from the highest sequence (`can_project`), the same rule as the launch, completion, and queue-failure writes. Chains after `0019_thread_incarnations` and is the current head.
- `persistence/bootstrap.py``bootstrap_schema(engine, backend=...)`, the three-branch provisioning decision, locked revision validation, and the narrow 0019 forward-compatibility exception - `persistence/bootstrap.py``bootstrap_schema(engine, backend=...)`, the three-branch provisioning decision, locked revision validation, and the narrow 0019 forward-compatibility exception
- `extensions/loader.py::load_extensions` — registers each spec's `table_prefix` with `register_extension_table_prefix()` - `extensions/loader.py::load_extensions` — registers each spec's `table_prefix` with `register_extension_table_prefix()`
- Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter, including extension-owned tables), `tests/test_extension_loader.py::TestTablePrefixRegistration` (spec-to-filter wiring), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor), `tests/test_migration_0004_run_ownership_dedupe.py` + `tests/test_migration_0007_scheduled_run_active_dedupe.py` (dedupe-before-unique-index pre-steps) - Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter, including extension-owned tables), `tests/test_extension_loader.py::TestTablePrefixRegistration` (spec-to-filter wiring), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor), `tests/test_migration_0004_run_ownership_dedupe.py` + `tests/test_migration_0007_scheduled_run_active_dedupe.py` (dedupe-before-unique-index pre-steps)

View File

@ -0,0 +1,43 @@
"""Persist per-task occurrence order and idempotent launch accounting.
Revision ID: 0022_scheduled_occurrence_seq
Revises: 0019_thread_incarnations
"""
from __future__ import annotations
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
revision: str = "0022_scheduled_occurrence_seq"
down_revision: str | Sequence[str] | None = "0019_thread_incarnations"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_INDEX = "uq_scheduled_task_run_occurrence_seq"
def upgrade() -> None:
from deerflow.persistence.migrations._helpers import safe_add_column
safe_add_column("scheduled_tasks", sa.Column("last_occurrence_seq", sa.BigInteger(), nullable=False, server_default="0"))
# Caller clocks cannot reconstruct the admission order or prove which
# historical launches were counted. Preserve NULL for both legacy fields.
safe_add_column("scheduled_task_runs", sa.Column("occurrence_seq", sa.BigInteger(), nullable=True))
safe_add_column("scheduled_task_runs", sa.Column("launch_accounted", sa.Boolean(), nullable=True))
indexes = {index["name"] for index in sa.inspect(op.get_bind()).get_indexes("scheduled_task_runs")}
if _INDEX not in indexes:
op.create_index(_INDEX, "scheduled_task_runs", ["task_id", "occurrence_seq"], unique=True)
def downgrade() -> None:
from deerflow.persistence.migrations._helpers import safe_drop_column
indexes = {index["name"] for index in sa.inspect(op.get_bind()).get_indexes("scheduled_task_runs")}
if _INDEX in indexes:
op.drop_index(_INDEX, table_name="scheduled_task_runs")
safe_drop_column("scheduled_task_runs", "launch_accounted")
safe_drop_column("scheduled_task_runs", "occurrence_seq")
safe_drop_column("scheduled_tasks", "last_occurrence_seq")

View File

@ -2,7 +2,7 @@ from __future__ import annotations
from datetime import UTC, datetime from datetime import UTC, datetime
from sqlalchemy import DateTime, Index, Integer, String, Text, text from sqlalchemy import BigInteger, Boolean, DateTime, Index, Integer, String, Text, text
from sqlalchemy.orm import Mapped, mapped_column from sqlalchemy.orm import Mapped, mapped_column
from deerflow.persistence.base import Base from deerflow.persistence.base import Base
@ -13,6 +13,10 @@ class ScheduledTaskRunRow(Base):
id: Mapped[str] = mapped_column(String(64), primary_key=True) id: Mapped[str] = mapped_column(String(64), primary_key=True)
task_id: Mapped[str] = mapped_column(String(64), index=True) task_id: Mapped[str] = mapped_column(String(64), index=True)
# NULL identifies legacy history or direct inserts without a parent task.
occurrence_seq: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
# New occurrences start False; NULL preserves unknown legacy accounting.
launch_accounted: Mapped[bool | None] = mapped_column(Boolean, nullable=True)
thread_id: Mapped[str] = mapped_column(String(64), index=True) thread_id: Mapped[str] = mapped_column(String(64), index=True)
run_id: Mapped[str | None] = mapped_column(String(64), nullable=True) run_id: Mapped[str | None] = mapped_column(String(64), nullable=True)
scheduled_for: Mapped[datetime] = mapped_column(DateTime(timezone=True)) scheduled_for: Mapped[datetime] = mapped_column(DateTime(timezone=True))
@ -27,6 +31,7 @@ class ScheduledTaskRunRow(Base):
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
__table_args__ = ( __table_args__ = (
Index("uq_scheduled_task_run_occurrence_seq", "task_id", "occurrence_seq", unique=True),
# At most one non-terminal (queued/launching/running) occurrence per # At most one non-terminal (queued/launching/running) occurrence per
# task. Queued occurrences are deliberately durable; ``launching`` is # task. Queued occurrences are deliberately durable; ``launching`` is
# a short lease-fenced claim used so multiple gateway instances cannot # a short lease-fenced claim used so multiple gateway instances cannot

View File

@ -0,0 +1,37 @@
"""Parent projection and launch accounting under the scheduled-task row lock."""
from __future__ import annotations
from typing import TYPE_CHECKING
from sqlalchemy.orm.attributes import flag_modified
if TYPE_CHECKING:
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow
def can_project(task: ScheduledTaskRow, occurrence: ScheduledTaskRunRow) -> bool:
"""Keep unsequenced history best-effort until a sequenced run is admitted."""
if occurrence.occurrence_seq is None:
return task.last_occurrence_seq == 0
return occurrence.occurrence_seq == task.last_occurrence_seq
def account_launch(task: ScheduledTaskRow, occurrence: ScheduledTaskRunRow, run_id: str) -> bool:
"""Count a proven launch once, in the same transaction as its marker.
Migrated NULL markers retain the old last_run_id inference for their first
repair; historical accounting cannot be reconstructed from occurrence times.
"""
if occurrence.launch_accounted is True:
return False
legacy = occurrence.launch_accounted is None
occurrence.launch_accounted = True
if legacy and task.last_run_id == run_id:
return False
task.run_count += 1
# A stale occurrence may change the count, but not the current projection's
# timestamp. Callers that also project a result explicitly set updated_at.
flag_modified(task, "updated_at")
return True

View File

@ -11,14 +11,12 @@ from sqlalchemy.orm import aliased
from deerflow.persistence.run import RunRepository from deerflow.persistence.run import RunRepository
from deerflow.persistence.run.model import RunRow from deerflow.persistence.run.model import RunRow
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow 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.scheduler.schedules import next_run_at as compute_next_run_at
from deerflow.utils.time import coerce_iso from deerflow.utils.time import coerce_iso
TERMINAL_RUN_STATUSES: frozenset[str] = frozenset({"success", "failed", "skipped", "interrupted"})
QUEUED_RUN_STATUSES: tuple[str, ...] = ("queued",)
EXECUTING_RUN_STATUSES: tuple[str, ...] = ("launching", "running") EXECUTING_RUN_STATUSES: tuple[str, ...] = ("launching", "running")
ACTIVE_RUN_STATUSES: tuple[str, ...] = (*QUEUED_RUN_STATUSES, *EXECUTING_RUN_STATUSES)
_SCHEDULER_BUDGET_LOCK_KEY = 4694001 _SCHEDULER_BUDGET_LOCK_KEY = 4694001
@ -65,7 +63,7 @@ class ScheduledTaskRunRepository:
@staticmethod @staticmethod
def _row_to_dict(row: ScheduledTaskRunRow) -> dict[str, Any]: def _row_to_dict(row: ScheduledTaskRunRow) -> dict[str, Any]:
data = row.to_dict() data = row.to_dict(exclude={"occurrence_seq", "launch_accounted"})
for key in ( for key in (
"scheduled_for", "scheduled_for",
"lease_expires_at", "lease_expires_at",
@ -103,7 +101,10 @@ class ScheduledTaskRunRepository:
candidate: RunRow, candidate: RunRow,
) -> None: ) -> None:
"""Repair the parent update if launch committed before bookkeeping.""" """Repair the parent update if launch committed before bookkeeping."""
if task is None or task.last_run_id == candidate.run_id: 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 return
launched_at = candidate.created_at launched_at = candidate.created_at
if launched_at.tzinfo is None: if launched_at.tzinfo is None:
@ -117,9 +118,9 @@ class ScheduledTaskRunRepository:
task.timezone, task.timezone,
now=launched_at, now=launched_at,
) )
task.run_count += 1
task.lease_owner = None task.lease_owner = None
task.lease_expires_at = None task.lease_expires_at = None
task.updated_at = datetime.now(UTC)
if task.schedule_type == "once": if task.schedule_type == "once":
if candidate.status == "success": if candidate.status == "success":
task.status = "completed" task.status = "completed"
@ -160,12 +161,14 @@ class ScheduledTaskRunRepository:
scheduled_for=scheduled_for, scheduled_for=scheduled_for,
trigger=trigger, trigger=trigger,
status=status, status=status,
launch_accounted=False,
created_at=datetime.now(UTC), created_at=datetime.now(UTC),
) )
async with self._sf() as session: async with self._sf() as session:
task: ScheduledTaskRow | None = None # 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 coordinate_with_task:
task = await self._lock_task(session, task_id)
if task is None or (expected_task_user_id is not None and task.user_id != expected_task_user_id): if task is None or (expected_task_user_id is not None and task.user_id != expected_task_user_id):
await session.rollback() await session.rollback()
raise ScheduledTaskAdmissionRejected(task_id, reason="not_found") raise ScheduledTaskAdmissionRejected(task_id, reason="not_found")
@ -191,8 +194,19 @@ class ScheduledTaskRunRepository:
if active_status is not None: if active_status is not None:
await session.rollback() await session.rollback()
raise ActiveScheduledRunConflict(task_id) 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) session.add(row)
if task is not None and release_task_lease_status is not None: if coordinate_with_task and task is not None and release_task_lease_status is not None:
task.status = release_task_lease_status task.status = release_task_lease_status
task.lease_owner = None task.lease_owner = None
task.lease_expires_at = None task.lease_expires_at = None
@ -201,12 +215,10 @@ class ScheduledTaskRunRepository:
await session.commit() await session.commit()
except IntegrityError: except IntegrityError:
await session.rollback() await session.rollback()
# Only active-status inserts can trip the partial unique index # A primary-key/sequence conflict is not necessarily an active
# ``uq_scheduled_task_run_active``; a terminal-status row (e.g. # slot conflict. Preserve the database error unless an active
# a "skipped" tombstone) is outside its predicate and cannot # occurrence actually exists after rollback.
# conflict, so any IntegrityError there is a genuine fault and 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)):
# is re-raised untranslated.
if status in ACTIVE_RUN_STATUSES:
raise ActiveScheduledRunConflict(task_id) from None raise ActiveScheduledRunConflict(task_id) from None
raise raise
await session.refresh(row) await session.refresh(row)
@ -412,7 +424,7 @@ class ScheduledTaskRunRepository:
row.lease_owner = None row.lease_owner = None
row.lease_expires_at = None row.lease_expires_at = None
if task is not None: if task is not None and can_project(task, row):
if task.status == "paused": if task.status == "paused":
# A concurrent/later pause owns the parent presentation. # A concurrent/later pause owns the parent presentation.
task.lease_owner = None task.lease_owner = None
@ -465,7 +477,7 @@ class ScheduledTaskRunRepository:
row.lease_owner = None row.lease_owner = None
row.lease_expires_at = None row.lease_expires_at = None
if task is not None: if task is not None and can_project(task, row):
if row.trigger == "manual": if row.trigger == "manual":
task_status = task.status or "enabled" task_status = task.status or "enabled"
next_at = task.next_run_at next_at = task.next_run_at

View File

@ -2,7 +2,7 @@ from __future__ import annotations
from datetime import UTC, datetime from datetime import UTC, datetime
from sqlalchemy import JSON, DateTime, Integer, String, Text from sqlalchemy import JSON, BigInteger, DateTime, Integer, String, Text
from sqlalchemy.orm import Mapped, mapped_column from sqlalchemy.orm import Mapped, mapped_column
from deerflow.persistence.base import Base from deerflow.persistence.base import Base
@ -31,9 +31,27 @@ class ScheduledTaskRow(Base):
lease_owner: Mapped[str | None] = mapped_column(String(128), nullable=True) lease_owner: Mapped[str | None] = mapped_column(String(128), nullable=True)
lease_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) lease_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
run_count: Mapped[int] = mapped_column(Integer, default=0) run_count: Mapped[int] = mapped_column(Integer, default=0)
last_occurrence_seq: Mapped[int] = mapped_column(BigInteger, default=0, server_default="0")
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
updated_at: Mapped[datetime] = mapped_column( updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), DateTime(timezone=True),
default=lambda: datetime.now(UTC), default=lambda: datetime.now(UTC),
onupdate=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC),
) )
# Status constants - shared between scheduled_tasks and scheduled_task_runs
# to avoid circular import and ensure consistency.
# Import these from deerflow.persistence.scheduled_tasks.model in both modules.
TERMINAL_RUN_STATUSES: frozenset[str] = frozenset({"success", "failed", "skipped", "interrupted"})
ACTIVE_RUN_STATUSES: frozenset[str] = frozenset({"queued", "launching", "running"})
# Parent ``once`` task status projected from a terminal occurrence status.
# Shared by the completion path and both recovery paths so the mapping
# cannot drift between them.
ONCE_TASK_STATUS_BY_RUN_STATUS: dict[str, str] = {
"success": "completed",
"failed": "failed",
"interrupted": "cancelled",
"skipped": "cancelled",
}

View File

@ -10,7 +10,9 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from deerflow.persistence.run import RunRepository from deerflow.persistence.run import RunRepository
from deerflow.persistence.run.model import RunRow from deerflow.persistence.run.model import RunRow
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow from deerflow.persistence.scheduled_task_runs.projection import account_launch, can_project
from deerflow.persistence.scheduled_tasks.model import ACTIVE_RUN_STATUSES, ONCE_TASK_STATUS_BY_RUN_STATUS, TERMINAL_RUN_STATUSES, ScheduledTaskRow
from deerflow.scheduler.schedules import next_run_at as compute_next_run_at
from deerflow.utils.time import coerce_iso from deerflow.utils.time import coerce_iso
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -64,7 +66,7 @@ class ScheduledTaskRepository:
@staticmethod @staticmethod
def _row_to_dict(row: ScheduledTaskRow) -> dict[str, Any]: def _row_to_dict(row: ScheduledTaskRow) -> dict[str, Any]:
data = row.to_dict() data = row.to_dict(exclude={"last_occurrence_seq"})
for key in ( for key in (
"created_at", "created_at",
"updated_at", "updated_at",
@ -364,17 +366,17 @@ class ScheduledTaskRepository:
async def release_queued_admission_lease(self, task_id: str) -> bool: async def release_queued_admission_lease(self, task_id: str) -> bool:
"""Recover a crash after queue insert but before parent-lease release.""" """Recover a crash after queue insert but before parent-lease release."""
async with self._sf() as session: async with self._sf() as session:
task = await session.get(ScheduledTaskRow, task_id, with_for_update=True) task = await self._lock_task(session, task_id)
if task is None or task.status != "running" or task.lease_owner is None: if task is None or task.status != "running" or task.lease_owner is None:
await session.rollback() await session.rollback()
return False return False
queued = await session.scalar( queued = await session.scalar(
select(ScheduledTaskRunRow.id).where( select(ScheduledTaskRunRow).where(
ScheduledTaskRunRow.task_id == task_id, ScheduledTaskRunRow.task_id == task_id,
ScheduledTaskRunRow.status == "queued", ScheduledTaskRunRow.status == "queued",
) )
) )
if queued is None: if queued is None or not can_project(task, queued):
await session.rollback() await session.rollback()
return False return False
task.status = "enabled" task.status = "enabled"
@ -397,9 +399,10 @@ class ScheduledTaskRepository:
increment_run_count: bool, increment_run_count: bool,
protect_terminal: bool = False, protect_terminal: bool = False,
expected_lease_owner: str | None = None, expected_lease_owner: str | None = None,
task_run_id: str | None = None,
) -> bool: ) -> bool:
async with self._sf() as session: async with self._sf() as session:
row = await session.get(ScheduledTaskRow, task_id, with_for_update=True) row = await self._lock_task(session, task_id)
if row is None: if row is None:
return False return False
if expected_lease_owner is not None and row.lease_owner != expected_lease_owner: if expected_lease_owner is not None and row.lease_owner != expected_lease_owner:
@ -411,16 +414,39 @@ class ScheduledTaskRepository:
) )
await session.rollback() await session.rollback()
return False return False
if protect_terminal and row.status in TERMINAL_TASK_STATUSES: occurrence = None
if task_run_id is not None:
occurrence = await session.get(ScheduledTaskRunRow, task_run_id, with_for_update=True)
if occurrence is None or occurrence.task_id != task_id or occurrence.run_id not in (None, last_run_id):
logger.warning(
"Fenced stale scheduled-task launch update for task %s: occurrence %s does not belong to run %s",
task_id,
task_run_id,
last_run_id,
)
await session.rollback()
return False
elif last_run_id is not None:
# Preserve direct repository callers that identify the launch
# by its durable run id rather than its occurrence id.
occurrence = await session.scalar(select(ScheduledTaskRunRow).where(ScheduledTaskRunRow.task_id == task_id, ScheduledTaskRunRow.run_id == last_run_id).with_for_update())
should_increment_run_count = increment_run_count and (last_run_id is None or row.last_run_id != last_run_id)
if occurrence is not None:
if increment_run_count and last_run_id is not None:
account_launch(row, occurrence, last_run_id)
should_increment_run_count = False
if not can_project(row, occurrence):
await session.commit()
return True
if protect_terminal and (row.status in TERMINAL_TASK_STATUSES or (occurrence is not None and occurrence.status in TERMINAL_RUN_STATUSES)):
# A fast-failing run can reach handle_run_completion (which # A fast-failing run can reach handle_run_completion (which
# finalizes a `once` task) before this launch-path write # finalizes a `once` task) before this launch-path write
# commits; keep the hook's status/error and only record the # commits. Cron parents stay enabled even after completion,
# launch bookkeeping. # so also protect the terminal occurrence's status/error.
pass pass
else: else:
row.status = status row.status = status
row.last_error = last_error row.last_error = last_error
should_increment_run_count = increment_run_count and (last_run_id is None or row.last_run_id != last_run_id)
row.next_run_at = _coerce_datetime(next_run_at) row.next_run_at = _coerce_datetime(next_run_at)
row.last_run_at = _coerce_datetime(last_run_at) row.last_run_at = _coerce_datetime(last_run_at)
row.last_run_id = last_run_id row.last_run_id = last_run_id
@ -433,6 +459,56 @@ class ScheduledTaskRepository:
await session.commit() await session.commit()
return True return True
async def complete_run(
self,
task_id: str,
*,
user_id: str,
task_run_id: str,
run_id: str,
status: str,
error: str | None,
finished_at: datetime,
) -> bool:
"""Commit occurrence completion, accounting and eligible parent outcome."""
if status not in TERMINAL_RUN_STATUSES:
raise ValueError(f"unsupported terminal occurrence status: {status!r}")
async with self._sf() as session:
task = await self._lock_task(session, task_id)
occurrence = await session.get(ScheduledTaskRunRow, task_run_id, with_for_update=True)
if occurrence is None or occurrence.task_id != task_id or occurrence.run_id not in (None, run_id) or (task is not None and task.user_id != user_id):
await session.rollback()
return False
occurrence.status = status
occurrence.run_id = run_id
occurrence.error = error
occurrence.finished_at = finished_at
occurrence.lease_owner = None
occurrence.lease_expires_at = None
if task is not None:
account_launch(task, occurrence, run_id)
if can_project(task, occurrence):
if task.last_run_id != run_id:
# A fast callback can beat launch bookkeeping. Finalize
# its association and schedule before releasing the slot.
launched_at = occurrence.started_at or occurrence.scheduled_for
if launched_at.tzinfo is None:
launched_at = launched_at.replace(tzinfo=UTC)
task.last_run_at = launched_at
task.last_run_id = run_id
task.last_thread_id = occurrence.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.last_error = error
if task.schedule_type == "once":
# Only a once task consumes its parent on completion;
# cron parents keep whatever status they already hold.
task.status = ONCE_TASK_STATUS_BY_RUN_STATUS[status]
task.updated_at = finished_at
await session.commit()
return True
async def claim_dispatch_lease( async def claim_dispatch_lease(
self, self,
task_id: str, task_id: str,
@ -477,6 +553,97 @@ class ScheduledTaskRepository:
result = await session.execute(stmt) result = await session.execute(stmt)
return [self._row_to_dict(row) for row in result.scalars()] return [self._row_to_dict(row) for row in result.scalars()]
@staticmethod
async def _fetch_latest_run(session: AsyncSession, task_id: str) -> ScheduledTaskRunRow | None:
"""Return the latest ``scheduled_task_runs`` row for a task (no status filter).
The caller must not hold a pre-lock snapshot of this row; the outcome
used for finalization must come from a fresh read. ``populate_existing``
bypasses the session identity map so a concurrently committed status is
read back fresh.
Once any sequenced row exists, the parent-locked ``occurrence_seq`` is
the only recency key and the highest sequence wins: caller clocks never
reorder sequenced rows, and unsequenced rows (legacy history or an
admission by a pre-upgrade writer) are not consulted, matching the
``can_project`` rule applied by every other parent write. Only a task
whose history is entirely unsequenced keeps the previous timestamp
ordering.
"""
base = select(ScheduledTaskRunRow).where(ScheduledTaskRunRow.task_id == task_id)
sequenced_stmt = base.where(ScheduledTaskRunRow.occurrence_seq.is_not(None)).order_by(ScheduledTaskRunRow.occurrence_seq.desc()).limit(1).execution_options(populate_existing=True)
sequenced = (await session.execute(sequenced_stmt)).scalars().first()
if sequenced is not None:
return sequenced
legacy_stmt = (
base.order_by(
ScheduledTaskRunRow.created_at.desc(),
ScheduledTaskRunRow.scheduled_for.desc(),
ScheduledTaskRunRow.id.desc(),
)
.limit(1)
.execution_options(populate_existing=True)
)
return (await session.execute(legacy_stmt)).scalars().first()
@staticmethod
async def _has_active_occurrence(session: AsyncSession, task_id: str) -> bool:
"""True while any occurrence row of the task is queued/launching/running.
``uq_scheduled_task_run_active`` allows one such row per task, so a live
row is the newest admission regardless of its caller clock or whether
it carries a sequence.
"""
stmt = (
select(ScheduledTaskRunRow.id)
.where(
ScheduledTaskRunRow.task_id == task_id,
ScheduledTaskRunRow.status.in_(ACTIVE_RUN_STATUSES),
)
.limit(1)
)
return (await session.execute(stmt)).scalars().first() is not None
@staticmethod
def _finalise_once_task_from_run(
task_row: ScheduledTaskRow,
run_row: ScheduledTaskRunRow | None,
*,
error: str,
now: datetime,
) -> bool:
"""Finalise a stuck ``once`` task parent to match its latest run outcome.
success -> completed, failed -> failed, interrupted -> cancelled,
skipped -> cancelled (no work performed). An active occurrence
(queued/launching/running) is left untouched a concurrent completion
or a later recovery pass will finalize it once the run reaches a
terminal state. When there is no run row at all the parent keeps
the original generic cancellation.
Returns ``True`` if the parent was finalised, ``False`` for an active
occurrence that must be retried by a later recovery pass.
"""
if run_row is not None and run_row.status in TERMINAL_RUN_STATUSES:
task_row.status = ONCE_TASK_STATUS_BY_RUN_STATUS[run_row.status]
if run_row.status == "success":
task_row.last_error = None
elif run_row.status == "interrupted":
task_row.last_error = run_row.error or error
else:
task_row.last_error = run_row.error
task_row.updated_at = now
return True
if run_row is not None and run_row.status in ACTIVE_RUN_STATUSES:
# Active occurrence — leave the parent unchanged. A concurrent
# completion or a later recovery pass will finalize it.
return False
# No run row, or an unrecognised legacy status — generic cancel.
task_row.status = "cancelled"
task_row.last_error = error
task_row.updated_at = now
return True
async def cancel_stuck_once_tasks(self, *, error: str) -> int: async def cancel_stuck_once_tasks(self, *, error: str) -> int:
"""Reconcile ``once`` tasks orphaned in ``running`` by a process crash. """Reconcile ``once`` tasks orphaned in ``running`` by a process crash.
@ -486,22 +653,51 @@ class ScheduledTaskRepository:
it. After a crash the hook is gone and the task would be stuck forever. it. After a crash the hook is gone and the task would be stuck forever.
Tasks still holding a lease are left alone they were claimed but not Tasks still holding a lease are left alone they were claimed but not
launched, and expired-lease reclaim recovers them safely. launched, and expired-lease reclaim recovers them safely.
Outcome-aware: for each stuck task, looks up the latest
``scheduled_task_runs`` row. If the run already reached a terminal
status the parent task is finalised to match (``success``
``completed``, ``failed`` ``failed``, ``interrupted``
``cancelled``, ``skipped`` ``cancelled``). While any occurrence
row is still active (queued/launching/running), sequenced or not, the
parent is left untouched: ``uq_scheduled_task_run_active`` makes that
row the task's newest admission. Tasks whose latest run row is absent
receive the generic cancellation.
""" """
stmt = select(ScheduledTaskRow).where( stmt = select(ScheduledTaskRow.id).where(
ScheduledTaskRow.schedule_type == "once", ScheduledTaskRow.schedule_type == "once",
ScheduledTaskRow.status == "running", ScheduledTaskRow.status == "running",
ScheduledTaskRow.lease_expires_at.is_(None), ScheduledTaskRow.lease_expires_at.is_(None),
) )
async with self._sf() as session: async with self._sf() as session:
result = await session.execute(stmt) task_ids = list((await session.execute(stmt)).scalars())
rows = list(result.scalars()) if not task_ids:
return 0
now = datetime.now(UTC) now = datetime.now(UTC)
for row in rows: reconciled = 0
row.status = "cancelled" for task_id in task_ids:
row.last_error = error # Row lock (no SQLite writer emulation, so the race regressions
row.updated_at = now # can still commit concurrently): on Postgres this serialises
# against admission, which locks the parent before inserting a
# queued occurrence, so no live row can appear between the
# probe below and this commit.
task_row = await session.get(ScheduledTaskRow, task_id, with_for_update=True)
if task_row is None or task_row.status != "running" or task_row.lease_expires_at is not None:
continue
run_row = await self._fetch_latest_run(session, task_id)
if await self._has_active_occurrence(session, task_id):
# The live row is the newest admission whatever its clock or
# sequence; its own completion, or a later pass once it is
# terminal, owns the parent.
continue
if run_row is not None and not can_project(task_row, run_row):
# Same eligibility rule as every other parent write: a row
# that cannot project leaves the parent untouched.
continue
if self._finalise_once_task_from_run(task_row, run_row, error=error, now=now):
reconciled += 1
await session.commit() await session.commit()
return len(rows) return reconciled
async def reconcile_stuck_once_tasks( async def reconcile_stuck_once_tasks(
self, self,
@ -540,9 +736,10 @@ class ScheduledTaskRepository:
if candidate is not None and candidate.status in {"pending", "running"}: 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): if _lease_is_alive(candidate.lease_expires_at, now=now, grace_seconds=lease_grace_seconds):
continue continue
# Run takeover commits in its own short transaction. If this # Run takeover commits the durable RunRow in its own short
# outer commit fails, the next poll finishes task bookkeeping # transaction. Its occurrence projection may remain active
# while the underlying run remains safely terminal. # until reconcile_active_runs runs, so this pass can defer the
# parent and let the next poll finish task bookkeeping.
claimed = await self._run_repository.claim_for_takeover( claimed = await self._run_repository.claim_for_takeover(
candidate.run_id, candidate.run_id,
grace_seconds=lease_grace_seconds, grace_seconds=lease_grace_seconds,
@ -553,10 +750,20 @@ class ScheduledTaskRepository:
refreshed = await self._run_repository.get(candidate.run_id, user_id=None) refreshed = await self._run_repository.get(candidate.run_id, user_id=None)
if refreshed is not None and refreshed.get("status") in {"pending", "running"}: if refreshed is not None and refreshed.get("status") in {"pending", "running"}:
continue continue
task.status = "cancelled" # Finalise from the latest scheduled_task_run (unconditional lookup).
task.last_error = error # Filtering by terminal status only could exclude a newer skipped
task.updated_at = datetime.now(UTC) # or active row, causing us to finalise based on an older run.
cancelled += 1 run_row = await self._fetch_latest_run(session, task.id)
if await self._has_active_occurrence(session, task.id):
# Any live occurrence, sequenced or not, is the newest
# admission; reconcile_active_runs terminalises it once its
# durable run is gone and the next pass finalises the parent.
continue
if run_row is not None and not can_project(task, run_row):
# Same eligibility rule as every other parent write.
continue
if self._finalise_once_task_from_run(task, run_row, error=error, now=now):
cancelled += 1
await session.commit() await session.commit()
return cancelled return cancelled

View File

@ -25,6 +25,58 @@ async def _noop_langgraph_runtime(_app, _startup_config):
yield yield
@asynccontextmanager
async def _langgraph_runtime_with_scheduler_repositories(app, _startup_config):
app.state.scheduled_task_repo = object()
app.state.scheduled_task_run_repo = object()
yield
def test_enabled_scheduler_start_failure_aborts_gateway_lifespan():
"""An enabled scheduler must fail lifespan before channel or request admission."""
from app.gateway.app import lifespan
async def scenario():
app = FastAPI()
startup_config = MagicMock()
startup_config.log_level = "INFO"
startup_config.memory.enabled = False
startup_config.memory.shutdown_flush_timeout_seconds = 5.0
startup_config.scheduler.enabled = True
startup_config.scheduler.multi_instance = False
startup_config.scheduler.poll_interval_seconds = 5
startup_config.scheduler.lease_seconds = 120
startup_config.scheduler.max_concurrent_runs = 3
startup_config.scheduler.queue_timeout_seconds = 3600
startup_config.run_ownership.grace_seconds = 10
channel_service = MagicMock()
channel_service.get_status.return_value = {}
start_channel_service = AsyncMock(return_value=channel_service)
scheduler_service = MagicMock()
scheduler_service.start = AsyncMock(side_effect=RuntimeError("scheduled recovery failed"))
scheduler_service.stop = AsyncMock()
with (
patch("app.gateway.app.get_app_config", return_value=startup_config),
patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)),
patch("app.gateway.app.langgraph_runtime", _langgraph_runtime_with_scheduler_repositories),
patch("app.gateway.app.auth.close_oidc_service", AsyncMock()),
patch("app.channels.service.start_channel_service", start_channel_service),
patch("app.channels.service.stop_channel_service", AsyncMock()),
patch("app.scheduler.ScheduledTaskService", return_value=scheduler_service),
patch("deerflow.skills.projection.ensure_public_skill_projection"),
patch("deerflow.agents.memory.get_memory_manager", return_value=MagicMock()),
):
with pytest.raises(RuntimeError, match="scheduled recovery failed"):
async with lifespan(app):
pass
scheduler_service.start.assert_awaited_once()
start_channel_service.assert_not_awaited()
asyncio.run(scenario())
async def _run_lifespan_with_hanging_stop() -> float: async def _run_lifespan_with_hanging_stop() -> float:
"""Drive the lifespan context with stop_channel_service hanging forever. """Drive the lifespan context with stop_channel_service hanging forever.

View File

@ -5,6 +5,7 @@ import asyncio
import pytest import pytest
import sqlalchemy as sa import sqlalchemy as sa
from alembic import command from alembic import command
from alembic.script import ScriptDirectory
from alembic.util.exc import CommandError from alembic.util.exc import CommandError
from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.ext.asyncio import create_async_engine
@ -65,15 +66,14 @@ async def test_forward_revision_cannot_skip_required_batch_columns(tmp_path, mon
await conn.execute(sa.text("ALTER TABLE mcp_tasks ADD COLUMN thread_incarnation VARCHAR(32)")) await conn.execute(sa.text("ALTER TABLE mcp_tasks ADD COLUMN thread_incarnation VARCHAR(32)"))
if race: if race:
current_head, current_revisions = bootstrap._get_revision_metadata() current_head, current_revisions = bootstrap._get_revision_metadata()
assert current_head == "0019_thread_incarnations" assert {"0020_threads_meta_project_id", "0021_batch_acceptance", "0019_thread_incarnations", current_head} <= current_revisions
assert {"0020_threads_meta_project_id", "0021_batch_acceptance", current_head} <= current_revisions # The published 0020 binary knows only the ancestors of its own head.
rollback_revisions = frozenset(revision.revision for revision in ScriptDirectory.from_config(cfg).iterate_revisions("0020_threads_meta_project_id", "base"))
assert not ({"0021_batch_acceptance", "0019_thread_incarnations", current_head} & rollback_revisions)
monkeypatch.setattr( monkeypatch.setattr(
bootstrap, bootstrap,
"_get_revision_metadata", "_get_revision_metadata",
lambda: ( lambda: ("0020_threads_meta_project_id", rollback_revisions),
"0020_threads_meta_project_id",
current_revisions - {"0021_batch_acceptance", current_head},
),
) )
def raced_upgrade(*args): def raced_upgrade(*args):

View File

@ -0,0 +1,159 @@
"""Upgrade existing scheduler history without guessing occurrence order or counts."""
from __future__ import annotations
import asyncio
import os
import uuid
from datetime import UTC, datetime
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import pytest
import pytest_asyncio
import sqlalchemy as sa
from alembic import command
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from deerflow.persistence.bootstrap import _get_alembic_config, _get_head_revision
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository
pytestmark = pytest.mark.asyncio
REVISION = "0022_scheduled_occurrence_seq"
PREVIOUS = "0021_batch_acceptance"
INDEX = "uq_scheduled_task_run_occurrence_seq"
@pytest_asyncio.fixture(params=["sqlite", "postgres"])
async def migration_database(request, tmp_path):
schema = None
if request.param == "postgres":
uri = os.environ.get("TEST_POSTGRES_URI")
if not uri:
pytest.skip("requires TEST_POSTGRES_URI (real Postgres migration)")
parts = urlsplit(uri)
# CI passes a sync ``postgresql://...?sslmode=disable`` URL; the async
# engine needs the asyncpg driver and rejects libpq-only query keys.
scheme = "postgresql+asyncpg" if parts.scheme in {"postgres", "postgresql"} else parts.scheme
query = urlencode([(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key not in {"sslmode", "channel_binding"}])
uri = urlunsplit(parts._replace(scheme=scheme, query=query))
schema = f"occurrence_migration_{uuid.uuid4().hex}"
engine = create_async_engine(uri, connect_args=build_asyncpg_connect_args(schema))
else:
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}")
cfg = _get_alembic_config(engine, postgres_schema=schema or "")
try:
if schema:
async with engine.begin() as connection:
await connection.execute(sa.text(f'CREATE SCHEMA "{schema}"'))
await asyncio.to_thread(command.upgrade, cfg, PREVIOUS)
async with engine.begin() as connection:
await connection.execute(
sa.text(
"INSERT INTO scheduled_tasks "
"(id, user_id, thread_id, context_mode, assistant_id, title, prompt, "
"schedule_type, schedule_spec, timezone, status, overlap_policy, run_count, created_at, updated_at) "
"VALUES ('task-legacy', 'user-1', 'thread-1', 'reuse_thread', 'lead_agent', "
"'Legacy', 'Prompt', 'cron', '{\"cron\":\"0 9 * * *\"}', 'UTC', "
"'enabled', 'enqueue', 7, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
)
)
for run_id in ("legacy-a", "legacy-b"):
await connection.execute(
sa.text("INSERT INTO scheduled_task_runs (id, task_id, thread_id, scheduled_for, trigger, status, created_at) VALUES (:run_id, 'task-legacy', 'thread-1', CURRENT_TIMESTAMP, 'manual', 'success', CURRENT_TIMESTAMP)"),
{"run_id": run_id},
)
yield engine, cfg
finally:
if schema:
async with engine.begin() as connection:
await connection.execute(sa.text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
await engine.dispose()
async def _schema(engine):
async with engine.connect() as connection:
return await connection.run_sync(
lambda conn: (
{column["name"]: column for column in sa.inspect(conn).get_columns("scheduled_tasks")},
{column["name"]: column for column in sa.inspect(conn).get_columns("scheduled_task_runs")},
{index["name"]: index for index in sa.inspect(conn).get_indexes("scheduled_task_runs")},
)
)
async def test_occurrence_revision_is_single_head():
assert _get_head_revision() == REVISION
async def test_upgrade_preserves_legacy_rows_and_allocates_from_one(migration_database):
engine, cfg = migration_database
before_task, before_run, _before_indexes = await _schema(engine)
assert "last_occurrence_seq" not in before_task
assert {"occurrence_seq", "launch_accounted"}.isdisjoint(before_run)
await asyncio.to_thread(command.upgrade, cfg, REVISION)
task_columns, run_columns, indexes = await _schema(engine)
assert task_columns["last_occurrence_seq"]["nullable"] is False
assert run_columns["occurrence_seq"]["nullable"] is True
assert run_columns["occurrence_seq"]["default"] is None
assert run_columns["launch_accounted"]["nullable"] is True
assert run_columns["launch_accounted"]["default"] is None
assert indexes[INDEX]["column_names"] == ["task_id", "occurrence_seq"]
assert indexes[INDEX]["unique"]
async with engine.connect() as connection:
legacy = (await connection.execute(sa.text("SELECT occurrence_seq, launch_accounted FROM scheduled_task_runs ORDER BY id"))).all()
task = (await connection.execute(sa.text("SELECT last_occurrence_seq, run_count, updated_at FROM scheduled_tasks WHERE id = 'task-legacy'"))).one()
assert legacy == [(None, None), (None, None)]
assert task.last_occurrence_seq == 0
assert task.run_count == 7
factory = async_sessionmaker(engine, expire_on_commit=False)
await ScheduledTaskRunRepository(factory).create(
run_record_id="new",
task_id="task-legacy",
thread_id="thread-new",
scheduled_for=datetime.now(UTC),
trigger="manual",
status="success",
)
async with engine.connect() as connection:
new_run = (await connection.execute(sa.text("SELECT occurrence_seq, launch_accounted FROM scheduled_task_runs WHERE id = 'new'"))).one()
after_task = (await connection.execute(sa.text("SELECT last_occurrence_seq, run_count, updated_at FROM scheduled_tasks WHERE id = 'task-legacy'"))).one()
assert new_run.occurrence_seq == 1
assert new_run.launch_accounted is not None
assert not new_run.launch_accounted
assert after_task.last_occurrence_seq == 1
assert after_task.run_count == 7
assert after_task.updated_at == task.updated_at
async def test_migration_retry_and_downgrade_preserve_history(migration_database):
engine, cfg = migration_database
await asyncio.to_thread(command.upgrade, cfg, REVISION)
# A retry can encounter the additive DDL already applied before stamping.
await asyncio.to_thread(command.stamp, cfg, PREVIOUS)
await asyncio.to_thread(command.upgrade, cfg, REVISION)
async with engine.connect() as connection:
assert await connection.scalar(sa.text("SELECT COUNT(*) FROM scheduled_task_runs")) == 2
assert await connection.scalar(sa.text("SELECT run_count FROM scheduled_tasks WHERE id = 'task-legacy'")) == 7
await asyncio.to_thread(command.downgrade, cfg, PREVIOUS)
task_columns, run_columns, indexes = await _schema(engine)
assert "last_occurrence_seq" not in task_columns
assert {"occurrence_seq", "launch_accounted"}.isdisjoint(run_columns)
assert INDEX not in indexes
await asyncio.to_thread(command.upgrade, cfg, REVISION)
async with engine.connect() as connection:
assert await connection.scalar(sa.text("SELECT COUNT(*) FROM scheduled_task_runs")) == 2
assert await connection.scalar(sa.text("SELECT run_count FROM scheduled_tasks WHERE id = 'task-legacy'")) == 7
async def test_migration_unique_index_accepts_nulls_and_rejects_duplicate_sequence(migration_database):
engine, cfg = migration_database
await asyncio.to_thread(command.upgrade, cfg, REVISION)
async with engine.begin() as connection:
await connection.execute(sa.text("UPDATE scheduled_task_runs SET occurrence_seq = 1 WHERE id = 'legacy-a'"))
with pytest.raises(sa.exc.IntegrityError):
async with connection.begin_nested():
await connection.execute(sa.text("UPDATE scheduled_task_runs SET occurrence_seq = 1 WHERE id = 'legacy-b'"))
await connection.execute(sa.text("UPDATE scheduled_task_runs SET occurrence_seq = 2 WHERE id = 'legacy-b'"))

View File

@ -13,6 +13,8 @@ from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import pytest import pytest
import sqlalchemy as sa import sqlalchemy as sa
from alembic import command as alembic_command from alembic import command as alembic_command
from alembic.config import Config as AlembicConfig
from alembic.script import ScriptDirectory
from alembic.util.exc import CommandError from alembic.util.exc import CommandError
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
@ -23,13 +25,15 @@ from deerflow.persistence.bootstrap import (
_CANONICAL_0019_SCHEMA_FLOOR, _CANONICAL_0019_SCHEMA_FLOOR,
_FORWARD_COMPATIBLE_REVISION, _FORWARD_COMPATIBLE_REVISION,
_get_alembic_config, _get_alembic_config,
_get_head_revision,
_upgrade, _upgrade,
bootstrap_schema, bootstrap_schema,
) )
from deerflow.persistence.engine import close_engine, get_engine, init_engine_from_config from deerflow.persistence.engine import close_engine, get_engine, init_engine_from_config
from deerflow.persistence.thread_meta.sql import ThreadMetaRepository from deerflow.persistence.thread_meta.sql import ThreadMetaRepository
CURRENT_HEAD = "0019_thread_incarnations" CANONICAL_INCARNATION_REVISION = "0019_thread_incarnations"
LOCAL_HEAD = _get_head_revision()
ROLLBACK_HEAD = "0020_threads_meta_project_id" ROLLBACK_HEAD = "0020_threads_meta_project_id"
INCARNATION_PARENT = "0021_batch_acceptance" INCARNATION_PARENT = "0021_batch_acceptance"
ORIGINAL_INCARNATION_PARENT = "0018_oauth_identity_pg_partial" ORIGINAL_INCARNATION_PARENT = "0018_oauth_identity_pg_partial"
@ -57,9 +61,18 @@ async def _set_database_revision(engine, revision: str) -> None:
await conn.execute(sa.text("UPDATE alembic_version SET version_num = :revision"), {"revision": revision}) await conn.execute(sa.text("UPDATE alembic_version SET version_num = :revision"), {"revision": revision})
async def _seed_current_head(engine) -> None: def _rollback_binary_revisions() -> frozenset[str]:
await bootstrap_schema(engine, backend="sqlite") """Revisions the published 0020 rollback binary knows: ancestors of its head."""
assert await _database_revision(engine) == CURRENT_HEAD cfg = AlembicConfig()
cfg.set_main_option("script_location", str(bootstrap_mod._MIGRATIONS_DIR))
script = ScriptDirectory.from_config(cfg)
return frozenset(revision.revision for revision in script.iterate_revisions(ROLLBACK_HEAD, "base"))
async def _seed_canonical_0019(engine) -> None:
"""Seed the exact canonical-0019 shape; the local head may already lie beyond it."""
await asyncio.to_thread(_upgrade, _get_alembic_config(engine), CANONICAL_INCARNATION_REVISION)
assert await _database_revision(engine) == CANONICAL_INCARNATION_REVISION
async def _seed_rollback_head(engine) -> None: async def _seed_rollback_head(engine) -> None:
@ -82,14 +95,16 @@ async def _add_forward_columns(engine) -> None:
def _simulate_rollback_binary(monkeypatch: pytest.MonkeyPatch) -> None: def _simulate_rollback_binary(monkeypatch: pytest.MonkeyPatch) -> None:
current_head, current_revisions = bootstrap_mod._get_revision_metadata() current_head, current_revisions = bootstrap_mod._get_revision_metadata()
assert current_head == CURRENT_HEAD == _FORWARD_COMPATIBLE_REVISION assert current_head == LOCAL_HEAD
assert ROLLBACK_HEAD in current_revisions assert CANONICAL_INCARNATION_REVISION == _FORWARD_COMPATIBLE_REVISION
assert INCARNATION_PARENT in current_revisions assert {ROLLBACK_HEAD, INCARNATION_PARENT, CANONICAL_INCARNATION_REVISION, LOCAL_HEAD} <= current_revisions
assert CURRENT_HEAD in current_revisions rollback_revisions = _rollback_binary_revisions()
assert ROLLBACK_HEAD in rollback_revisions
assert not ({INCARNATION_PARENT, CANONICAL_INCARNATION_REVISION, LOCAL_HEAD} & rollback_revisions)
monkeypatch.setattr( monkeypatch.setattr(
bootstrap_mod, bootstrap_mod,
"_get_revision_metadata", "_get_revision_metadata",
lambda: (ROLLBACK_HEAD, current_revisions - {INCARNATION_PARENT, CURRENT_HEAD}), lambda: (ROLLBACK_HEAD, rollback_revisions),
) )
@ -104,7 +119,7 @@ async def _seed_original_forward_schema(engine) -> None:
async def test_canonical_0019_floor_matches_migration_schema(tmp_path: Path) -> None: async def test_canonical_0019_floor_matches_migration_schema(tmp_path: Path) -> None:
engine = create_async_engine(_url(tmp_path, "canonical-floor.db")) engine = create_async_engine(_url(tmp_path, "canonical-floor.db"))
try: try:
await asyncio.to_thread(_upgrade, _get_alembic_config(engine), CURRENT_HEAD) await asyncio.to_thread(_upgrade, _get_alembic_config(engine), CANONICAL_INCARNATION_REVISION)
async with engine.connect() as conn: async with engine.connect() as conn:
def reflect(sync_conn): def reflect(sync_conn):
@ -191,7 +206,7 @@ async def test_audited_original_forward_schema_can_upgrade_preserving_incarnatio
await asyncio.to_thread(alembic_command.stamp, _get_alembic_config(engine), ORIGINAL_INCARNATION_PARENT, purge=True) await asyncio.to_thread(alembic_command.stamp, _get_alembic_config(engine), ORIGINAL_INCARNATION_PARENT, purge=True)
await bootstrap_schema(engine, backend="sqlite") await bootstrap_schema(engine, backend="sqlite")
assert await _database_revision(engine) == CURRENT_HEAD assert await _database_revision(engine) == LOCAL_HEAD
repository = ThreadMetaRepository(async_sessionmaker(engine, expire_on_commit=False)) repository = ThreadMetaRepository(async_sessionmaker(engine, expire_on_commit=False))
assert [row["thread_id"] for row in await repository.search(user_id=None)] == ["existing"] assert [row["thread_id"] for row in await repository.search(user_id=None)] == ["existing"]
assert (await repository.create("new", user_id=None))["thread_id"] == "new" assert (await repository.create("new", user_id=None))["thread_id"] == "new"
@ -214,7 +229,7 @@ async def test_known_canonical_0019_validates_fixed_floor_then_upgrades_to_futur
future_table = None future_table = None
calls: list[str] = [] calls: list[str] = []
try: try:
await _seed_current_head(engine) await _seed_canonical_0019(engine)
# Model a future binary whose ORM includes schema that only its next # Model a future binary whose ORM includes schema that only its next
# migration can add. Canonical 0019 must not be rejected for lacking it. # migration can add. Canonical 0019 must not be rejected for lacking it.
future_table = sa.Table("future_after_0019", Base.metadata, sa.Column("id", sa.String(), primary_key=True)) future_table = sa.Table("future_after_0019", Base.metadata, sa.Column("id", sa.String(), primary_key=True))
@ -273,7 +288,7 @@ async def test_known_older_revision_upgrades_normally(tmp_path: Path) -> None:
await bootstrap_schema(engine, backend="sqlite") await bootstrap_schema(engine, backend="sqlite")
assert await _database_revision(engine) == CURRENT_HEAD assert await _database_revision(engine) == LOCAL_HEAD
finally: finally:
await engine.dispose() await engine.dispose()
@ -286,7 +301,7 @@ async def test_exact_forward_revision_skips_upgrade_with_warning(
) -> None: ) -> None:
engine = create_async_engine(_url(tmp_path, "forward.db")) engine = create_async_engine(_url(tmp_path, "forward.db"))
try: try:
await _seed_current_head(engine) await _seed_canonical_0019(engine)
_simulate_rollback_binary(monkeypatch) _simulate_rollback_binary(monkeypatch)
with caplog.at_level("WARNING", logger="deerflow.persistence.bootstrap"): with caplog.at_level("WARNING", logger="deerflow.persistence.bootstrap"):
@ -302,7 +317,7 @@ async def test_exact_forward_revision_skips_upgrade_with_warning(
async def test_other_unknown_revision_fails_closed(tmp_path: Path) -> None: async def test_other_unknown_revision_fails_closed(tmp_path: Path) -> None:
engine = create_async_engine(_url(tmp_path, "unknown.db")) engine = create_async_engine(_url(tmp_path, "unknown.db"))
try: try:
await _seed_current_head(engine) await _seed_canonical_0019(engine)
await _set_database_revision(engine, "9999_unknown") await _set_database_revision(engine, "9999_unknown")
with pytest.raises(RuntimeError, match="not known to this build"): with pytest.raises(RuntimeError, match="not known to this build"):
@ -328,7 +343,7 @@ async def test_sqlite_upgrade_race_recovers_when_other_process_applies_forward_r
upgrade_started.set() upgrade_started.set()
if not continue_upgrade.wait(timeout=5): if not continue_upgrade.wait(timeout=5):
raise TimeoutError("timed out waiting for the forward migration") raise TimeoutError("timed out waiting for the forward migration")
raise CommandError(f"Can't locate revision identified by '{CURRENT_HEAD}'") raise CommandError(f"Can't locate revision identified by '{CANONICAL_INCARNATION_REVISION}'")
try: try:
await _seed_rollback_head(old_gateway) await _seed_rollback_head(old_gateway)
@ -339,7 +354,7 @@ async def test_sqlite_upgrade_race_recovers_when_other_process_applies_forward_r
assert await asyncio.to_thread(upgrade_started.wait, 5) assert await asyncio.to_thread(upgrade_started.wait, 5)
new_cfg = _get_alembic_config(new_gateway) new_cfg = _get_alembic_config(new_gateway)
await asyncio.to_thread(_upgrade, new_cfg, CURRENT_HEAD) await asyncio.to_thread(_upgrade, new_cfg, CANONICAL_INCARNATION_REVISION)
with caplog.at_level("WARNING", logger="deerflow.persistence.bootstrap"): with caplog.at_level("WARNING", logger="deerflow.persistence.bootstrap"):
continue_upgrade.set() continue_upgrade.set()
@ -395,7 +410,7 @@ async def test_local_forward_migration_error_stays_fatal(
async def test_empty_alembic_version_fails_closed(tmp_path: Path) -> None: async def test_empty_alembic_version_fails_closed(tmp_path: Path) -> None:
engine = create_async_engine(_url(tmp_path, "empty-version.db")) engine = create_async_engine(_url(tmp_path, "empty-version.db"))
try: try:
await _seed_current_head(engine) await _seed_canonical_0019(engine)
async with engine.begin() as conn: async with engine.begin() as conn:
await conn.execute(sa.text("DELETE FROM alembic_version")) await conn.execute(sa.text("DELETE FROM alembic_version"))
@ -409,7 +424,7 @@ async def test_empty_alembic_version_fails_closed(tmp_path: Path) -> None:
async def test_multiple_alembic_versions_fail_closed(tmp_path: Path) -> None: async def test_multiple_alembic_versions_fail_closed(tmp_path: Path) -> None:
engine = create_async_engine(_url(tmp_path, "multiple-versions.db")) engine = create_async_engine(_url(tmp_path, "multiple-versions.db"))
try: try:
await _seed_current_head(engine) await _seed_canonical_0019(engine)
async with engine.begin() as conn: async with engine.begin() as conn:
await conn.execute( await conn.execute(
sa.text("INSERT INTO alembic_version (version_num) VALUES (:revision)"), sa.text("INSERT INTO alembic_version (version_num) VALUES (:revision)"),
@ -426,7 +441,7 @@ async def test_multiple_alembic_versions_fail_closed(tmp_path: Path) -> None:
async def test_rollback_batch_writer_tolerates_acceptance_columns(tmp_path: Path) -> None: async def test_rollback_batch_writer_tolerates_acceptance_columns(tmp_path: Path) -> None:
engine = create_async_engine(_url(tmp_path, "batch-repository.db")) engine = create_async_engine(_url(tmp_path, "batch-repository.db"))
try: try:
await _seed_current_head(engine) await _seed_canonical_0019(engine)
old_items = sa.table( old_items = sa.table(
"subagent_batch_items", "subagent_batch_items",
sa.column("id"), sa.column("id"),
@ -485,7 +500,7 @@ async def test_rollback_batch_writer_tolerates_acceptance_columns(tmp_path: Path
async def test_rollback_thread_writer_tolerates_forward_nullable_column(tmp_path: Path) -> None: async def test_rollback_thread_writer_tolerates_forward_nullable_column(tmp_path: Path) -> None:
engine = create_async_engine(_url(tmp_path, "thread-repository.db")) engine = create_async_engine(_url(tmp_path, "thread-repository.db"))
try: try:
await _seed_current_head(engine) await _seed_canonical_0019(engine)
# This is the complete 0020 table shape. Keeping it independent from # This is the complete 0020 table shape. Keeping it independent from
# the current ORM prevents a future model change from silently making # the current ORM prevents a future model change from silently making
# this rollback-writer test aware of the forward column. # this rollback-writer test aware of the forward column.
@ -543,7 +558,7 @@ async def test_rollback_thread_writer_tolerates_forward_nullable_column(tmp_path
async def test_rollback_shaped_mcp_task_sql_tolerates_forward_nullable_column(tmp_path: Path) -> None: async def test_rollback_shaped_mcp_task_sql_tolerates_forward_nullable_column(tmp_path: Path) -> None:
engine = create_async_engine(_url(tmp_path, "mcp-repository.db")) engine = create_async_engine(_url(tmp_path, "mcp-repository.db"))
try: try:
await _seed_current_head(engine) await _seed_canonical_0019(engine)
# This is the complete 0020 task table shape, deliberately excluding # This is the complete 0020 task table shape, deliberately excluding
# only the forward thread_incarnation column. # only the forward thread_incarnation column.
old_tasks = sa.table( old_tasks = sa.table(
@ -690,7 +705,11 @@ async def test_old_gateway_restarts_against_forward_postgres_revision(
await init_engine_from_config(config) await init_engine_from_config(config)
engine = get_engine() engine = get_engine()
assert engine is not None assert engine is not None
assert await _database_revision(engine) == CURRENT_HEAD assert await _database_revision(engine) == LOCAL_HEAD
# The 0020 rollback binary allowlists canonical 0019 only, so model the
# audited window by stepping the schema back to that exact revision.
await asyncio.to_thread(alembic_command.downgrade, _get_alembic_config(engine, postgres_schema=schema), CANONICAL_INCARNATION_REVISION)
assert await _database_revision(engine) == CANONICAL_INCARNATION_REVISION
await close_engine() await close_engine()
_simulate_rollback_binary(monkeypatch) _simulate_rollback_binary(monkeypatch)

View File

@ -0,0 +1,314 @@
"""Database ordering is per task and survives retries independently of caller clocks."""
from __future__ import annotations
import asyncio
import os
import uuid
from datetime import UTC, datetime, timedelta
from unittest.mock import patch
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import pytest
import pytest_asyncio
from sqlalchemy import select, text
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
import deerflow.persistence.models # noqa: F401
from deerflow.persistence.base import Base
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
from deerflow.persistence.scheduled_task_runs import ActiveScheduledRunConflict, ScheduledTaskRunRepository
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
from deerflow.persistence.scheduled_tasks.model import ACTIVE_RUN_STATUSES, ONCE_TASK_STATUS_BY_RUN_STATUS, ScheduledTaskRow
pytestmark = pytest.mark.asyncio
@pytest_asyncio.fixture(params=["sqlite", "postgres"])
async def occurrence_factories(request, tmp_path):
"""Two pools guarantee competing admissions use independent DB connections."""
schema = None
if request.param == "postgres":
uri = os.environ.get("TEST_POSTGRES_URI")
if not uri:
pytest.skip("requires TEST_POSTGRES_URI (real Postgres for occurrence ordering)")
parts = urlsplit(uri)
# CI passes a sync ``postgresql://...?sslmode=disable`` URL; the async
# engine needs the asyncpg driver and rejects libpq-only query keys.
scheme = "postgresql+asyncpg" if parts.scheme in {"postgres", "postgresql"} else parts.scheme
query = urlencode([(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key not in {"sslmode", "channel_binding"}])
uri = urlunsplit(parts._replace(scheme=scheme, query=query))
schema = f"occurrence_{uuid.uuid4().hex}"
options = {"connect_args": build_asyncpg_connect_args(schema)}
else:
uri = f"sqlite+aiosqlite:///{tmp_path / 'occurrences.db'}"
options = {"connect_args": {"timeout": 30}}
engines = [create_async_engine(uri, **options) for _ in range(2)]
try:
async with engines[0].begin() as connection:
if schema:
await connection.execute(text(f'CREATE SCHEMA "{schema}"'))
await connection.run_sync(Base.metadata.create_all)
yield tuple(async_sessionmaker(engine, expire_on_commit=False) for engine in engines)
finally:
if schema:
async with engines[0].begin() as connection:
await connection.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
for engine in engines:
await engine.dispose()
async def _create_task(factory, task_id="task", *, schedule_type="cron"):
spec = {"cron": "* * * * *"} if schedule_type == "cron" else {"run_at": datetime(2026, 7, 15, 12, 0, tzinfo=UTC).isoformat()}
return await ScheduledTaskRepository(factory).create(
task_id=task_id,
user_id="user-1",
thread_id="thread-1",
context_mode="reuse_thread",
assistant_id=None,
title="Occurrence ordering",
prompt="p",
schedule_type=schedule_type,
schedule_spec=spec,
timezone="UTC",
next_run_at=None,
)
async def _create_run(factory, run_id, *, task_id="task", status="success"):
return await ScheduledTaskRunRepository(factory).create(
run_record_id=run_id,
task_id=task_id,
thread_id=f"thread-{run_id}",
scheduled_for=datetime(2026, 7, 15, 12, 0, tzinfo=UTC),
trigger="manual",
status=status,
)
async def _sequence(factory, run_id):
async with factory() as session:
return await session.scalar(select(ScheduledTaskRunRow.occurrence_seq).where(ScheduledTaskRunRow.id == run_id))
async def _high_water_mark(factory, task_id="task"):
async with factory() as session:
return await session.scalar(select(ScheduledTaskRow.last_occurrence_seq).where(ScheduledTaskRow.id == task_id))
async def test_concurrent_allocations_use_distinct_monotonic_sequences(occurrence_factories):
first, second = occurrence_factories
original = await _create_task(first)
ready = [asyncio.Event(), asyncio.Event()]
start = asyncio.Event()
async def admit(factory, lane):
ready[lane].set()
await start.wait()
for index in range(4):
await _create_run(factory, f"run-{lane}-{index}")
admissions = [asyncio.create_task(admit(factory, lane)) for lane, factory in enumerate((first, second))]
await asyncio.gather(*(event.wait() for event in ready))
start.set()
await asyncio.gather(*admissions)
sequences = [await _sequence(first, f"run-{lane}-{index}") for lane in range(2) for index in range(4)]
assert sorted(sequences) == list(range(1, 9))
for lane in range(2):
lane_sequences = sequences[lane * 4 : (lane + 1) * 4]
assert lane_sequences == sorted(lane_sequences)
assert await _high_water_mark(first) == 8
current = await ScheduledTaskRepository(first).get("task", user_id="user-1")
assert current["updated_at"] == original["updated_at"]
assert current["run_count"] == 0
async def test_sequence_allocation_is_independent_per_task(occurrence_factories):
first, second = occurrence_factories
for task_id in ("task-a", "task-b"):
await _create_task(first, task_id)
await _create_run(first, "run-a1", task_id="task-a")
await _create_run(second, "run-a2", task_id="task-a")
await _create_run(second, "run-b1", task_id="task-b")
assert [await _sequence(first, run_id) for run_id in ("run-a1", "run-a2", "run-b1")] == [1, 2, 1]
async def test_active_conflict_rolls_back_sequence_allocation(occurrence_factories):
first, second = occurrence_factories
await _create_task(first)
await _create_run(first, "active", status="queued")
with pytest.raises(ActiveScheduledRunConflict):
await _create_run(second, "rejected", status="queued")
assert await _high_water_mark(first) == 1
assert await _sequence(first, "rejected") is None
await ScheduledTaskRunRepository(first).update_status("active", status="success")
await _create_run(second, "accepted", status="queued")
assert await _sequence(first, "accepted") == 2
@pytest.mark.parametrize("status", ["queued", "success"])
async def test_primary_key_conflict_is_not_an_active_conflict_and_rolls_back(occurrence_factories, status):
first, second = occurrence_factories
await _create_task(first, "task-a")
await _create_task(first, "task-b")
await _create_run(first, "duplicate", task_id="task-a")
with pytest.raises(IntegrityError):
await _create_run(second, "duplicate", task_id="task-b", status=status)
assert await _high_water_mark(first, "task-b") == 0
await _create_run(second, "unique", task_id="task-b", status=status)
assert await _sequence(first, "unique") == 1
async def test_requeue_and_reclaim_preserve_occurrence_sequence(occurrence_factories):
first, second = occurrence_factories
await _create_task(first)
await _create_run(first, "retry", status="queued")
now = datetime(2026, 7, 15, 12, 0, tzinfo=UTC)
repo = ScheduledTaskRunRepository(second)
for attempt in range(2):
claimed = await repo.claim_queued_run("retry", now=now, lease_owner="worker", lease_seconds=60, global_max_concurrent_runs=1)
assert claimed is not None
assert claimed["attempt_count"] == attempt + 1
assert await repo.requeue_claimed_run("retry", lease_owner="worker") is True
assert await _sequence(first, "retry") == 1
assert await _high_water_mark(first) == 1
async def test_internal_sequence_fields_are_absent_from_repository_responses(occurrence_factories):
first, _second = occurrence_factories
created_task = await _create_task(first)
task_repo = ScheduledTaskRepository(first)
run_repo = ScheduledTaskRunRepository(first)
created_run = await _create_run(first, "queued", status="queued")
task_responses = [created_task, await task_repo.get("task", user_id="user-1"), *(await task_repo.list_by_user("user-1"))]
run_responses = [created_run, await run_repo.get_active_run("task"), *(await run_repo.list_by_task("task")), *(await run_repo.list_queued_runs(limit=10))]
for response in task_responses + run_responses:
assert {"last_occurrence_seq", "occurrence_seq", "launch_accounted"}.isdisjoint(response)
assert await _sequence(first, "queued") == 1
async def _insert_unsequenced_run(factory, run_id, *, created_at, status="success"):
"""Insert without the repository: legacy history or a pre-upgrade writer."""
async with factory() as session:
session.add(
ScheduledTaskRunRow(
id=run_id,
task_id="task",
thread_id=f"thread-{run_id}",
scheduled_for=created_at,
created_at=created_at,
trigger="manual",
status=status,
)
)
await session.commit()
async def _latest_run_id(factory):
async with factory() as session:
latest = await ScheduledTaskRepository._fetch_latest_run(session, "task")
assert latest is not None
return latest.id
async def test_recovery_order_keeps_timestamp_fallback_for_legacy_only_history(occurrence_factories):
first, _second = occurrence_factories
await _create_task(first)
now = datetime(2026, 7, 15, 12, 0, tzinfo=UTC)
# Unsequenced history is deliberately not assigned guessed sequence values.
for index in range(2):
await _insert_unsequenced_run(first, f"legacy-{index}", created_at=now + timedelta(days=index))
assert await _latest_run_id(first) == "legacy-1"
assert await _sequence(first, "legacy-0") is None
assert await _sequence(first, "legacy-1") is None
assert await _high_water_mark(first) == 0
@pytest.mark.parametrize("unsequenced_status", ["skipped", "running"])
@pytest.mark.parametrize("unsequenced_offset", [timedelta(days=-365), timedelta(seconds=30)], ids=["legacy-history-is-older", "pre-upgrade-writer-is-newer"])
async def test_recovery_lookup_prefers_the_highest_sequence_whenever_one_exists(occurrence_factories, unsequenced_offset, unsequenced_status):
"""Sequence decides whenever a sequenced row exists.
Reversed caller clocks between sequenced rows do not matter, and an
unsequenced row (legacy history or a pre-upgrade Gateway writer) is not
consulted even when its caller timestamp is later: the lookup returns the
same row ``can_project`` accepts, so recovery cannot act on a row that the
other parent writes would reject.
"""
first, second = occurrence_factories
await _create_task(first)
now = datetime(2026, 7, 15, 12, 0, tzinfo=UTC)
# The later sequence carries the earlier caller clock: sequence still wins.
with patch("deerflow.persistence.scheduled_task_runs.sql.datetime") as clock:
clock.now.return_value = now + timedelta(seconds=30)
await _create_run(first, "older-sequenced")
clock.now.return_value = now
await _create_run(second, "newer-sequenced")
await _insert_unsequenced_run(first, "unsequenced", created_at=now + unsequenced_offset, status=unsequenced_status)
assert [await _sequence(first, run_id) for run_id in ("older-sequenced", "newer-sequenced", "unsequenced")] == [1, 2, None]
assert await _latest_run_id(first) == "newer-sequenced"
@pytest.mark.parametrize("recovery_method", ["cancel_stuck_once_tasks", "reconcile_stuck_once_tasks"])
@pytest.mark.parametrize(
("sequenced_status", "unsequenced_status"),
[("running", "skipped"), ("success", "running"), ("success", "interrupted"), ("failed", "running")],
ids=[
"unsequenced-skipped-while-sequenced-active",
"unsequenced-running-while-sequenced-success",
"unsequenced-interrupted-while-sequenced-success",
"unsequenced-running-while-sequenced-failed",
],
)
async def test_once_recovery_defers_while_any_occurrence_is_live_then_projects_the_sequence_winner(occurrence_factories, recovery_method, sequenced_status, unsequenced_status):
"""Mixed-writer interleavings from review, both recovery paths, both backends.
A live occurrence row is the task's newest admission by construction
(``uq_scheduled_task_run_active``), whatever its caller clock and whether or
not it carries a sequence, so recovery defers while one exists: an
unsequenced ``skipped`` row cannot cancel a parent whose sequenced
occurrence is live, and an unsequenced ``running`` row cannot be skipped
over to finalise the parent from an older sequenced outcome. Once no row is
live, the sequence winner decides and a terminalised unsequenced row never
overrides it.
"""
first, _second = occurrence_factories
await _create_task(first, schedule_type="once")
task_repo = ScheduledTaskRepository(first)
await task_repo.update("task", user_id="user-1", updates={"status": "running"})
await _create_run(first, "sequenced", status=sequenced_status)
# A pre-upgrade node on a skewed clock: no sequence, later caller timestamp.
await _insert_unsequenced_run(first, "unsequenced", created_at=datetime.now(UTC) + timedelta(minutes=5), status=unsequenced_status)
assert await _sequence(first, "sequenced") == 1
assert await _high_water_mark(first) == 1
kwargs = {"error": "interrupted: recovery"}
if recovery_method == "reconcile_stuck_once_tasks":
kwargs["now"] = datetime.now(UTC) + timedelta(minutes=10)
async def recover():
count = await getattr(task_repo, recovery_method)(**kwargs)
task = await task_repo.get_internal("task")
assert task is not None
return count, task
any_live = sequenced_status in ACTIVE_RUN_STATUSES or unsequenced_status in ACTIVE_RUN_STATUSES
sequence_outcome = ONCE_TASK_STATUS_BY_RUN_STATUS.get(sequenced_status, "running")
expected_status, expected_count = ("running", 0) if any_live else (sequence_outcome, 1)
for _ in range(2): # a second pass must not change the outcome
count, task = await recover()
assert task["status"] == expected_status
assert task["last_error"] is None
assert count == expected_count
expected_count = 0
if unsequenced_status in ACTIVE_RUN_STATUSES:
# The pre-upgrade node died and occurrence recovery terminalised its
# row: the sequence winner now decides, not the newer unsequenced row.
assert await ScheduledTaskRunRepository(first).update_status("unsequenced", status="interrupted", error="pre-upgrade node died")
count, task = await recover()
assert task["status"] == sequence_outcome
assert task["last_error"] is None
assert count == 1

View File

@ -4,6 +4,7 @@ import asyncio
import os import os
import uuid import uuid
from datetime import UTC, datetime, timedelta from datetime import UTC, datetime, timedelta
from unittest.mock import patch
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import pytest import pytest
@ -14,6 +15,7 @@ from deerflow.config.database_config import DatabaseConfig
from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config
from deerflow.persistence.run import RunRepository from deerflow.persistence.run import RunRepository
from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
POSTGRES_URL = os.environ.get("TEST_POSTGRES_URI") POSTGRES_URL = os.environ.get("TEST_POSTGRES_URI")
@ -156,3 +158,71 @@ async def test_postgres_reconciliation_uses_metadata_and_atomically_claims_expir
assert recovered is not None assert recovered is not None
assert recovered["status"] == "error" assert recovered["status"] == "error"
assert recovered["stop_reason"] == "scheduled_task_orphan_recovered" assert recovered["stop_reason"] == "scheduled_task_orphan_recovered"
@pytest.mark.asyncio
@pytest.mark.parametrize("recovery_method", ["cancel_stuck_once_tasks", "reconcile_stuck_once_tasks"])
@pytest.mark.parametrize("newer_status", ["skipped", "queued", "launching", "running"])
@pytest.mark.parametrize("older_clock_ahead_seconds", [30, 0], ids=["reversed-timestamps", "equal-timestamps"])
async def test_postgres_once_recovery_uses_occurrence_order_despite_clock_skew(postgres_repositories, recovery_method, newer_status, older_clock_ahead_seconds):
task_repo, task_run_repo, _run_repo = postgres_repositories
now = datetime(2026, 7, 15, 12, 0, tzinfo=UTC)
await task_repo.create(
task_id="task-once",
user_id="user-1",
thread_id=None,
context_mode="fresh_thread_per_run",
assistant_id=None,
title="Once task",
prompt="p",
schedule_type="once",
schedule_spec={"run_at": now.isoformat()},
timezone="UTC",
next_run_at=None,
)
await task_repo.update_after_launch(
"task-once",
status="running",
next_run_at=None,
last_run_at=now,
last_run_id=None,
last_thread_id=None,
last_error=None,
increment_run_count=False,
)
# The older worker's clock and descending UUID order both favor its success.
with patch("deerflow.persistence.scheduled_task_runs.sql.datetime") as clock:
clock.now.return_value = now + timedelta(seconds=older_clock_ahead_seconds)
older = await task_run_repo.create(
run_record_id="ffffffff-ffff-4fff-8fff-ffffffffffff",
task_id="task-once",
thread_id="thread-old",
scheduled_for=clock.now.return_value,
trigger="manual",
status="success",
)
clock.now.return_value = now
newer = await task_run_repo.create(
run_record_id="00000000-0000-4000-8000-000000000000",
task_id="task-once",
thread_id="thread-new",
scheduled_for=clock.now.return_value,
trigger="manual",
status=newer_status,
)
assert older["created_at"] >= newer["created_at"]
assert older["scheduled_for"] >= newer["scheduled_for"]
kwargs = {"error": "interrupted: recovery"}
if recovery_method == "reconcile_stuck_once_tasks":
kwargs["now"] = now + timedelta(minutes=1)
count = await getattr(task_repo, recovery_method)(**kwargs)
task = await task_repo.get_internal("task-once")
assert task is not None
assert task["status"] == ("cancelled" if newer_status == "skipped" else "running")
assert count == (1 if newer_status == "skipped" else 0)
assert task["last_error"] is None
async with task_run_repo._sf() as session:
older_row = await session.get(ScheduledTaskRunRow, older["id"])
newer_row = await session.get(ScheduledTaskRunRow, newer["id"])
assert newer_row.occurrence_seq > older_row.occurrence_seq

View File

@ -154,6 +154,94 @@ async def test_queued_run_survives_single_instance_restart_sweep(tmp_path):
await close_engine() await close_engine()
async def test_queued_once_task_survives_startup_and_is_drained_on_next_poll(tmp_path):
"""A newer queued occurrence survives recovery of an already-stuck parent.
The old success models a completion that committed before its parent update.
A real manual dispatch then leaves newer work queued after a transient
same-thread conflict. Startup must not let the old success finalize the
parent through the newer active row; the ordinary queue drain owns launch.
"""
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
try:
sf = get_session_factory()
assert sf is not None
task_repo = ScheduledTaskRepository(sf)
run_repo = ScheduledTaskRunRepository(sf)
now = datetime.now(UTC)
await task_repo.create(
task_id="task-queued-once",
user_id="user-1",
thread_id="thread-queued-once",
context_mode="reuse_thread",
assistant_id="lead_agent",
title="Queued once task",
prompt="Resume queued work",
schedule_type="once",
schedule_spec={"run_at": now.isoformat()},
timezone="UTC",
next_run_at=None,
)
await task_repo.update(
"task-queued-once",
user_id="user-1",
updates={"status": "running"},
)
await run_repo.create(
run_record_id="task-run-old-success",
task_id="task-queued-once",
thread_id="thread-queued-once",
scheduled_for=now - timedelta(minutes=1),
trigger="scheduled",
status="success",
)
task = await task_repo.get("task-queued-once", user_id="user-1")
assert task is not None
launched = []
async def launch_run(**kwargs):
launched.append(kwargs)
if len(launched) == 1:
raise ConflictError("Thread thread-queued-once already has an active run")
return {"run_id": "run-after-restart", "thread_id": kwargs["thread_id"]}
first_service = _make_service(task_repo, run_repo, launch_run)
queued = await first_service.dispatch_task(task, now=now, trigger="manual")
assert queued["outcome"] == "queued"
rows = await run_repo.list_by_task("task-queued-once")
assert [row["status"] for row in rows] == ["queued", "success"]
service = _make_service(task_repo, run_repo, launch_run)
async def parked_run_loop():
await service._stop.wait()
service._run_loop = parked_run_loop
await service.start()
try:
task = await task_repo.get_internal("task-queued-once")
rows = await run_repo.list_by_task("task-queued-once")
assert task is not None
assert task["status"] == "running"
assert rows[0]["status"] == "queued"
await service.run_once(now=now + timedelta(seconds=1))
task = await task_repo.get_internal("task-queued-once")
rows = await run_repo.list_by_task("task-queued-once")
assert task is not None
assert task["status"] == "running"
assert task["last_run_id"] == "run-after-restart"
assert rows[0]["status"] == "running"
assert rows[0]["run_id"] == "run-after-restart"
assert rows[1]["status"] == "success"
assert len(launched) == 2
finally:
await service.stop()
finally:
await close_engine()
async def test_only_one_worker_can_claim_a_queued_run(tmp_path): async def test_only_one_worker_can_claim_a_queued_run(tmp_path):
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
try: try:

View File

@ -14,6 +14,7 @@ class DummyTaskRepo:
self.rows = rows self.rows = rows
self.claimed = False self.claimed = False
self.updated = None self.updated = None
self.completions = []
self.release_calls = [] self.release_calls = []
self.cancelled_stuck_once = None self.cancelled_stuck_once = None
self.reconciled_stuck_once = None self.reconciled_stuck_once = None
@ -45,6 +46,10 @@ class DummyTaskRepo:
async def update_after_launch(self, *args, **kwargs): async def update_after_launch(self, *args, **kwargs):
self.updated = (args, kwargs) self.updated = (args, kwargs)
async def complete_run(self, task_id, **kwargs):
self.completions.append((task_id, kwargs))
return True
async def get(self, task_id: str, *, user_id: str): async def get(self, task_id: str, *, user_id: str):
row = next((item for item in self.rows if item["id"] == task_id and item["user_id"] == user_id), None) row = next((item for item in self.rows if item["id"] == task_id and item["user_id"] == user_id), None)
return dict(row) if row is not None else None return dict(row) if row is not None else None
@ -452,7 +457,7 @@ async def test_run_once_continues_batch_after_invalid_thread_id():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_handle_run_completion_persists_success(): async def test_handle_run_completion_uses_atomic_repository_boundary():
task_repo = DummyTaskRepo( task_repo = DummyTaskRepo(
[ [
{ {
@ -494,9 +499,16 @@ async def test_handle_run_completion_persists_success():
await service.handle_run_completion(record) await service.handle_run_completion(record)
assert run_repo.updated[-1][0] == "task-run-6" assert len(task_repo.completions) == 1
assert run_repo.updated[-1][1]["status"] == "success" task_id, completion = task_repo.completions[0]
assert task_repo.rows[0]["last_error"] is None assert task_id == "task-6"
assert completion["user_id"] == "user-1"
assert completion["task_run_id"] == "task-run-6"
assert completion["run_id"] == "run-6"
assert completion["status"] == "success"
assert completion["error"] is None
assert completion["finished_at"].tzinfo == UTC
assert run_repo.updated == []
def _make_service(task_repo, run_repo): def _make_service(task_repo, run_repo):
@ -542,57 +554,29 @@ def _completion_record(status, *, task_id="task-once", error=None):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_once_task_completes_only_via_completion_hook(): @pytest.mark.parametrize(
("run_status", "error", "occurrence_status", "expected_error"),
[
(RunStatus.success, None, "success", None),
(RunStatus.error, "boom", "failed", "boom"),
(RunStatus.timeout, "time limit", "failed", "time limit"),
(RunStatus.interrupted, None, "interrupted", "run was interrupted before completion"),
(RunStatus.interrupted, "cancelled by user", "interrupted", "cancelled by user"),
],
)
async def test_handle_run_completion_forwards_terminal_outcome(run_status, error, occurrence_status, expected_error):
task_repo = DummyTaskRepo([_once_task_row()]) task_repo = DummyTaskRepo([_once_task_row()])
run_repo = DummyRunRepo() run_repo = DummyRunRepo()
service = _make_service(task_repo, run_repo) service = _make_service(task_repo, run_repo)
await service.handle_run_completion(_completion_record(RunStatus.success)) await service.handle_run_completion(_completion_record(run_status, error=error))
assert run_repo.updated[-1][1]["status"] == "success" assert len(task_repo.completions) == 1
assert task_repo.rows[0]["status"] == "completed" task_id, completion = task_repo.completions[0]
assert task_id == "task-once"
assert completion["status"] == occurrence_status
@pytest.mark.asyncio assert completion["error"] == expected_error
async def test_once_task_failed_run_marks_task_failed(): assert run_repo.updated == []
task_repo = DummyTaskRepo([_once_task_row()])
run_repo = DummyRunRepo()
service = _make_service(task_repo, run_repo)
await service.handle_run_completion(_completion_record(RunStatus.error, error="boom"))
assert run_repo.updated[-1][1]["status"] == "failed"
assert run_repo.updated[-1][1]["error"] == "boom"
assert task_repo.rows[0]["status"] == "failed"
assert task_repo.rows[0]["last_error"] == "boom"
@pytest.mark.asyncio
async def test_interrupted_run_is_distinct_and_cancels_once_task():
task_repo = DummyTaskRepo([_once_task_row()])
run_repo = DummyRunRepo()
service = _make_service(task_repo, run_repo)
await service.handle_run_completion(_completion_record(RunStatus.interrupted))
run_update = run_repo.updated[-1][1]
assert run_update["status"] == "interrupted"
assert run_update["error"] == "run was interrupted before completion"
assert task_repo.rows[0]["status"] == "cancelled"
@pytest.mark.asyncio
async def test_interrupted_cron_run_keeps_task_enabled():
row = _once_task_row(task_id="task-cron")
row.update({"schedule_type": "cron", "schedule_spec": {"cron": "0 9 * * *"}, "status": "enabled"})
task_repo = DummyTaskRepo([row])
run_repo = DummyRunRepo()
service = _make_service(task_repo, run_repo)
await service.handle_run_completion(_completion_record(RunStatus.interrupted, task_id="task-cron"))
assert run_repo.updated[-1][1]["status"] == "interrupted"
assert task_repo.rows[0]["status"] == "enabled"
@pytest.mark.asyncio @pytest.mark.asyncio
@ -636,6 +620,64 @@ async def test_startup_sweep_reconciles_stale_runs_and_stuck_once_tasks():
assert task_repo.cancelled_stuck_once == run_repo.stale_marked assert task_repo.cancelled_stuck_once == run_repo.stale_marked
@pytest.mark.parametrize("failure_stage", ["occurrence", "parent"])
@pytest.mark.asyncio
async def test_single_instance_start_fails_closed_before_polling(failure_stage):
order = []
class StartupTaskRepo(DummyTaskRepo):
def __init__(self):
super().__init__([])
self.recovery_attempts = 0
async def cancel_stuck_once_tasks(self, *, error):
self.recovery_attempts += 1
order.append("parent")
assert service._task is None
if failure_stage == "parent" and self.recovery_attempts == 1:
raise RuntimeError("simulated parent recovery failure")
return await super().cancel_stuck_once_tasks(error=error)
class StartupRunRepo(DummyRunRepo):
def __init__(self):
super().__init__()
self.recovery_attempts = 0
async def mark_stale_active_runs(self, *, error):
self.recovery_attempts += 1
order.append("occurrence")
assert service._task is None
if failure_stage == "occurrence" and self.recovery_attempts == 1:
raise RuntimeError("simulated occurrence recovery failure")
return await super().mark_stale_active_runs(error=error)
task_repo = StartupTaskRepo()
run_repo = StartupRunRepo()
service = ScheduledTaskService(
task_repo=task_repo,
task_run_repo=run_repo,
launch_run=lambda **_kwargs: None,
poll_interval_seconds=0,
lease_seconds=120,
max_concurrent_runs=3,
)
async def parked_run_loop():
await service._stop.wait()
service._run_loop = parked_run_loop
try:
with pytest.raises(RuntimeError, match=f"simulated {failure_stage} recovery failure"):
await service.start()
assert run_repo.recovery_attempts == 1
assert task_repo.recovery_attempts == (0 if failure_stage == "occurrence" else 1)
assert order == (["occurrence"] if failure_stage == "occurrence" else ["occurrence", "parent"])
assert service._task is None
assert task_repo.claimed is False
finally:
await service.stop()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_multi_instance_start_uses_lease_aware_reconciliation(): async def test_multi_instance_start_uses_lease_aware_reconciliation():
task_repo = DummyTaskRepo([]) task_repo = DummyTaskRepo([])

View File

@ -0,0 +1,187 @@
from contextlib import asynccontextmanager
from datetime import UTC, datetime, timedelta
import pytest
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from app.scheduler.service import ScheduledTaskService
from deerflow.persistence.base import Base
from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow
from deerflow.runtime import RunStatus
from deerflow.runtime.runs.manager import RunRecord
from deerflow.runtime.runs.schemas import DisconnectMode
@asynccontextmanager
async def _database(tmp_path):
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'completion.db'}")
try:
async with engine.begin() as connection:
await connection.run_sync(lambda sync: Base.metadata.create_all(sync, tables=[ScheduledTaskRow.__table__, ScheduledTaskRunRow.__table__]))
yield engine, async_sessionmaker(engine, expire_on_commit=False)
finally:
await engine.dispose()
async def _launching_occurrence(sf, *, schedule_type):
now = datetime(2026, 9, 9, 8, tzinfo=UTC)
tasks = ScheduledTaskRepository(sf)
occurrences = ScheduledTaskRunRepository(sf)
spec = {"run_at": now.isoformat()} if schedule_type == "once" else {"cron": "0 9 * * *"}
await tasks.create(
task_id="task-completion",
user_id="user-1",
thread_id=None,
context_mode="fresh_thread_per_run",
assistant_id=None,
title="completion",
prompt="test",
schedule_type=schedule_type,
schedule_spec=spec,
timezone="UTC",
next_run_at=now,
)
await occurrences.create(
run_record_id="occurrence-completion",
task_id="task-completion",
thread_id="thread-completion",
scheduled_for=now,
trigger="manual",
status="queued",
)
claimed = await occurrences.claim_queued_run(
"occurrence-completion",
lease_owner="launcher",
now=now,
lease_seconds=120,
global_max_concurrent_runs=3,
)
assert claimed is not None
await tasks.update(
"task-completion",
user_id="user-1",
updates={"status": "running" if schedule_type == "once" else "enabled", "last_error": "previous error"},
)
return tasks, occurrences, now
def _service(tasks, occurrences):
return ScheduledTaskService(
task_repo=tasks,
task_run_repo=occurrences,
launch_run=lambda **_kwargs: None,
poll_interval_seconds=5,
lease_seconds=120,
max_concurrent_runs=3,
)
def _completion(status, error):
return RunRecord(
run_id="run-completion",
thread_id="thread-completion",
assistant_id=None,
status=status,
on_disconnect=DisconnectMode.continue_,
metadata={"scheduled_task_id": "task-completion", "scheduled_task_run_id": "occurrence-completion"},
user_id="user-1",
error=error,
)
async def _snapshot(sf):
async with sf() as session:
task = await session.get(ScheduledTaskRow, "task-completion")
occurrence = await session.get(ScheduledTaskRunRow, "occurrence-completion")
assert task is not None and occurrence is not None
return task.to_dict(), occurrence.to_dict()
@pytest.mark.asyncio
@pytest.mark.parametrize("schedule_type", ["once", "cron"])
@pytest.mark.parametrize(
("run_status", "error", "occurrence_status", "once_status", "expected_error"),
[
(RunStatus.success, None, "success", "completed", None),
(RunStatus.error, "boom", "failed", "failed", "boom"),
(RunStatus.timeout, "time limit", "failed", "failed", "time limit"),
(RunStatus.interrupted, None, "interrupted", "cancelled", "run was interrupted before completion"),
],
)
async def test_completion_before_launch_bookkeeping_is_durable_and_counted_once(tmp_path, schedule_type, run_status, error, occurrence_status, once_status, expected_error):
async with _database(tmp_path) as (_engine, sf):
tasks, occurrences, now = await _launching_occurrence(sf, schedule_type=schedule_type)
service = _service(tasks, occurrences)
record = _completion(run_status, error)
await service.handle_run_completion(record)
task, occurrence = await _snapshot(sf)
assert task["status"] == (once_status if schedule_type == "once" else "enabled")
assert task["last_error"] == expected_error
assert task["last_run_id"] == "run-completion"
assert task["last_thread_id"] == "thread-completion"
assert task["run_count"] == 1
assert occurrence["status"] == occurrence_status
assert occurrence["error"] == expected_error
assert occurrence["run_id"] == "run-completion"
assert occurrence["finished_at"] is not None
assert occurrence["launch_accounted"] is True
assert occurrence["lease_owner"] is None
assert occurrence["lease_expires_at"] is None
# The late launcher and a repeated delivery must not undo the terminal
# outcome or count the same launch a second time.
await tasks.update_after_launch(
"task-completion",
task_run_id="occurrence-completion",
status="running" if schedule_type == "once" else "enabled",
next_run_at=None if schedule_type == "once" else now + timedelta(hours=1),
last_run_at=now,
last_run_id="run-completion",
last_thread_id="thread-completion",
last_error=None,
increment_run_count=True,
protect_terminal=True,
)
task_after_launch, _ = await _snapshot(sf)
assert task_after_launch["run_count"] == 1
assert task_after_launch["status"] == task["status"]
assert task_after_launch["last_error"] == expected_error
await service.handle_run_completion(record)
task_after_retry, _ = await _snapshot(sf)
assert task_after_retry["run_count"] == 1
assert task_after_retry["status"] == task["status"]
assert task_after_retry["last_error"] == expected_error
@pytest.mark.asyncio
async def test_completion_commit_failure_rolls_back_outcome_and_accounting(tmp_path):
async with _database(tmp_path) as (engine, sf):
tasks, occurrences, _now = await _launching_occurrence(sf, schedule_type="once")
before = await _snapshot(sf)
class FailingCommitSession(AsyncSession):
async def commit(self):
await self.flush()
raise RuntimeError("completion commit failed")
failing_sf = async_sessionmaker(engine, class_=FailingCommitSession, expire_on_commit=False)
failing_service = _service(ScheduledTaskRepository(failing_sf), occurrences)
record = _completion(RunStatus.success, None)
with pytest.raises(RuntimeError, match="completion commit failed"):
await failing_service.handle_run_completion(record)
# In particular, failure cannot leave a terminal child that active-run
# recovery would skip while the parent count still lacks this launch.
assert await _snapshot(sf) == before
service = _service(tasks, occurrences)
await service.handle_run_completion(record)
await service.handle_run_completion(record)
task, occurrence = await _snapshot(sf)
assert task["status"] == "completed"
assert task["run_count"] == 1
assert occurrence["status"] == "success"
assert occurrence["launch_accounted"] is True

View File

@ -0,0 +1,947 @@
"""Tests for scheduler completion consistency fixes.
Bug 1: handle_run_completion crash between two transactions mislabels
successful once-tasks as cancelled.
Bug 2: restart reconciliation (cancel_stuck_once_tasks and
reconcile_stuck_once_tasks) was not outcome-aware.
These tests exercise the REAL repositories against file-backed SQLite so
that partial-unique-index and ORM constraints are enforced.
"""
from __future__ import annotations
from datetime import UTC, datetime, timedelta
from unittest.mock import patch
import pytest
from sqlalchemy import update
from app.scheduler.service import ScheduledTaskService
from deerflow.config.database_config import DatabaseConfig
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow
from deerflow.runtime.runs.manager import RunRecord
from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus
pytestmark = pytest.mark.asyncio
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
_NOW = datetime(2026, 7, 15, 12, 0, 0, tzinfo=UTC)
async def _init_db(tmp_path) -> tuple[ScheduledTaskRepository, ScheduledTaskRunRepository]:
"""Set up a fresh file-backed SQLite database and return both repos."""
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
sf = get_session_factory()
assert sf is not None
return ScheduledTaskRepository(sf), ScheduledTaskRunRepository(sf)
async def _create_once_task(
task_repo: ScheduledTaskRepository,
*,
task_id: str = "task-once-1",
) -> dict:
"""Insert a once-type task (status set by the caller via ``_set_task_running``)."""
return await task_repo.create(
task_id=task_id,
user_id="user-1",
thread_id=None,
context_mode="fresh_thread_per_run",
assistant_id=None,
title="Once Task",
prompt="do it",
schedule_type="once",
schedule_spec={"run_at": "2026-07-15T12:00:00Z"},
timezone="UTC",
next_run_at=None,
)
async def _create_cron_task(
task_repo: ScheduledTaskRepository,
*,
task_id: str = "task-cron-1",
) -> dict:
"""Insert a cron-type task."""
now = _NOW
return await task_repo.create(
task_id=task_id,
user_id="user-1",
thread_id=None,
context_mode="fresh_thread_per_run",
assistant_id=None,
title="Cron Task",
prompt="do it daily",
schedule_type="cron",
schedule_spec={"cron": "0 9 * * *"},
timezone="UTC",
next_run_at=now + timedelta(hours=1),
)
async def _create_run(
run_repo: ScheduledTaskRunRepository,
*,
run_id: str = "task-run-1",
task_id: str = "task-once-1",
status: str = "running",
error: str | None = None,
created_at: datetime | None = None,
scheduled_for: datetime | None = None,
) -> dict:
"""Insert a scheduled_task_runs row."""
now = created_at or _NOW
run_row = await run_repo.create(
run_record_id=run_id,
task_id=task_id,
thread_id="thread-1",
scheduled_for=scheduled_for or now,
trigger="scheduled",
status=status,
)
# Explicitly set created_at so multi-history regressions exercise the
# primary recency key behind _fetch_latest_run. scheduled_for still
# defaults to the logical occurrence time unless a test overrides it.
if created_at is not None:
async with run_repo._sf() as session:
await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == run_id).values(created_at=created_at))
await session.commit()
return run_row
async def _set_task_running(task_repo: ScheduledTaskRepository, task_id: str) -> None:
"""Directly set a task to 'running' status (simulating post-launch state)."""
sf = get_session_factory()
assert sf is not None
async with sf() as session:
await session.execute(update(ScheduledTaskRow).where(ScheduledTaskRow.id == task_id).values(status="running", lease_owner=None, lease_expires_at=None, updated_at=_NOW))
await session.commit()
async def _get_task(task_repo: ScheduledTaskRepository, task_id: str) -> dict | None:
"""Read a task dict, bypassing user_id check for internal access."""
return await task_repo.get_internal(task_id)
# ---------------------------------------------------------------------------
# Bug 1 & 2: cancel_stuck_once_tasks outcome-aware recovery
# ---------------------------------------------------------------------------
class TestCancelStuckOnceTasksOutcomeAware:
"""Verify that cancel_stuck_once_tasks respects terminal run statuses."""
async def test_successful_run_reconciles_to_completed(self, tmp_path):
"""Crash after run committed 'success' but before parent update.
cancel_stuck_once_tasks must mark the parent 'completed', not
'cancelled'.
"""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-success")
await _set_task_running(task_repo, "task-success")
# Run row is terminal success — the completion hook wrote it
# but crashed before updating the parent task.
await _create_run(run_repo, run_id="run-success", task_id="task-success", status="success")
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 1
task = await _get_task(task_repo, "task-success")
assert task is not None
assert task["status"] == "completed"
assert task["last_error"] is None
finally:
await close_engine()
async def test_failed_run_reconciles_to_failed_with_error(self, tmp_path):
"""Crash after run committed 'failed' but before parent update.
cancel_stuck_once_tasks must mark the parent 'failed' and
propagate the run's own error message.
"""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-failed")
await _set_task_running(task_repo, "task-failed")
await _create_run(run_repo, run_id="run-failed", task_id="task-failed", status="failed")
# Write the error on the run row.
sf = get_session_factory()
assert sf is not None
async with sf() as session:
await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == "run-failed").values(error="LLM rate limited"))
await session.commit()
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 1
task = await _get_task(task_repo, "task-failed")
assert task is not None
assert task["status"] == "failed"
assert task["last_error"] == "LLM rate limited"
finally:
await close_engine()
async def test_interrupted_run_reconciles_to_cancelled_with_run_error(self, tmp_path):
"""Crash after run committed 'interrupted' but before parent update.
cancel_stuck_once_tasks must mark the parent 'cancelled' and use
the run's own error if present.
"""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-interrupted")
await _set_task_running(task_repo, "task-interrupted")
await _create_run(run_repo, run_id="run-int", task_id="task-interrupted", status="interrupted")
sf = get_session_factory()
assert sf is not None
async with sf() as session:
await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == "run-int").values(error="user cancelled"))
await session.commit()
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 1
task = await _get_task(task_repo, "task-interrupted")
assert task is not None
assert task["status"] == "cancelled"
assert task["last_error"] == "user cancelled"
finally:
await close_engine()
async def test_interrupted_run_without_error_uses_recovery_message(self, tmp_path):
"""Interrupted run with no error of its own should use the recovery error."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-int-norun")
await _set_task_running(task_repo, "task-int-norun")
await _create_run(run_repo, run_id="run-int2", task_id="task-int-norun", status="interrupted")
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 1
task = await _get_task(task_repo, "task-int-norun")
assert task is not None
assert task["status"] == "cancelled"
assert task["last_error"] == "interrupted: gateway restarted"
finally:
await close_engine()
async def test_no_run_row_preserves_original_cancel_behaviour(self, tmp_path):
"""No run row at all → generic cancel (active runs are left untouched)."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-norun")
await _set_task_running(task_repo, "task-norun")
# No run row created at all.
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 1
task = await _get_task(task_repo, "task-norun")
assert task is not None
assert task["status"] == "cancelled"
assert task["last_error"] == "interrupted: gateway restarted"
finally:
await close_engine()
async def test_active_run_leaves_parent_unchanged(self, tmp_path):
"""An active (non-terminal) run row → parent left unchanged, not cancelled."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-active-run")
await _set_task_running(task_repo, "task-active-run")
await _create_run(run_repo, run_id="run-active", task_id="task-active-run", status="running")
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 0
task = await _get_task(task_repo, "task-active-run")
assert task is not None
assert task["status"] == "running"
finally:
await close_engine()
async def test_recurring_task_not_affected(self, tmp_path):
"""Recurring tasks with status='running' must not be touched."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_cron_task(task_repo, task_id="task-cron")
# Manually set it to running (simulating an in-progress cron task).
sf = get_session_factory()
assert sf is not None
async with sf() as session:
await session.execute(update(ScheduledTaskRow).where(ScheduledTaskRow.id == "task-cron").values(status="running", lease_owner=None, lease_expires_at=None))
await session.commit()
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
# cancel_stuck_once_tasks only selects schedule_type == "once".
assert count == 0
task = await _get_task(task_repo, "task-cron")
assert task is not None
assert task["status"] == "running"
finally:
await close_engine()
async def test_leased_once_task_not_cancelled(self, tmp_path):
"""A once task still holding a lease is left alone."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-leased")
# Set status=running but keep a future lease.
sf = get_session_factory()
assert sf is not None
future = _NOW + timedelta(hours=1)
async with sf() as session:
await session.execute(
update(ScheduledTaskRow)
.where(ScheduledTaskRow.id == "task-leased")
.values(
status="running",
lease_owner="worker-1",
lease_expires_at=future,
)
)
await session.commit()
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 0
task = await _get_task(task_repo, "task-leased")
assert task is not None
assert task["status"] == "running"
finally:
await close_engine()
async def test_multiple_stuck_tasks_mixed_outcomes(self, tmp_path):
"""Multiple stuck tasks with different run statuses are reconciled correctly."""
task_repo, run_repo = await _init_db(tmp_path)
try:
# Task 1: successful run
await _create_once_task(task_repo, task_id="task-mixed-1")
await _set_task_running(task_repo, "task-mixed-1")
await _create_run(run_repo, run_id="run-m1", task_id="task-mixed-1", status="success")
# Task 2: failed run with error
await _create_once_task(task_repo, task_id="task-mixed-2")
await _set_task_running(task_repo, "task-mixed-2")
await _create_run(run_repo, run_id="run-m2", task_id="task-mixed-2", status="failed")
sf = get_session_factory()
assert sf is not None
async with sf() as session:
await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == "run-m2").values(error="timeout"))
await session.commit()
# Task 3: no run row
await _create_once_task(task_repo, task_id="task-mixed-3")
await _set_task_running(task_repo, "task-mixed-3")
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 3
t1 = await _get_task(task_repo, "task-mixed-1")
assert t1["status"] == "completed"
assert t1["last_error"] is None
t2 = await _get_task(task_repo, "task-mixed-2")
assert t2["status"] == "failed"
assert t2["last_error"] == "timeout"
t3 = await _get_task(task_repo, "task-mixed-3")
assert t3["status"] == "cancelled"
assert t3["last_error"] == "interrupted: gateway restarted"
finally:
await close_engine()
async def test_cancel_sees_run_committed_before_finalize(self, tmp_path):
"""Monkeypatch _fetch_latest_run to simulate the race: a completion
commits success after recovery has observed the stale 'running' state.
The test proves the post-lock fresh read picks up the committed success.
A reverted pre-lock batch implementation never calls _fetch_latest_run,
so the monkeypatch never fires and the test fails catching the regression.
"""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-cancel-race")
await _set_task_running(task_repo, "task-cancel-race")
await _create_run(run_repo, run_id="run-crace", task_id="task-cancel-race", status="running")
original_fetch = ScheduledTaskRepository._fetch_latest_run
fresh_sf = get_session_factory()
assert fresh_sf is not None
intercepted = False
async def _intercepted_fetch(session, task_id: str):
nonlocal intercepted
if not intercepted and task_id == "task-cancel-race":
intercepted = True
# Concurrent completion commits success in a separate session.
async with fresh_sf() as cs:
await cs.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == "run-crace").values(status="success"))
await cs.commit()
# Fresh read with populate_existing — production code's guarantee.
return await original_fetch(session, task_id)
return await original_fetch(session, task_id)
ScheduledTaskRepository._fetch_latest_run = staticmethod(_intercepted_fetch)
try:
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 1
task = await _get_task(task_repo, "task-cancel-race")
assert task is not None
assert task["status"] == "completed"
assert task["last_error"] is None
finally:
ScheduledTaskRepository._fetch_latest_run = staticmethod(original_fetch)
finally:
await close_engine()
# ---------------------------------------------------------------------------
# Bug 1 & 2: reconcile_stuck_once_tasks (multi-instance path)
# ---------------------------------------------------------------------------
class TestReconcileStuckOnceTasksOutcomeAware:
"""Verify the multi-instance path is also outcome-aware."""
async def test_successful_run_reconciles_to_completed(self, tmp_path):
"""reconcile_stuck_once_tasks with terminal success run → completed."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-recon-success")
await _set_task_running(task_repo, "task-recon-success")
await _create_run(run_repo, run_id="run-rs1", task_id="task-recon-success", status="success")
count = await task_repo.reconcile_stuck_once_tasks(
error="interrupted: lease expired",
now=_NOW,
lease_grace_seconds=10,
)
assert count == 1
task = await _get_task(task_repo, "task-recon-success")
assert task is not None
assert task["status"] == "completed"
assert task["last_error"] is None
finally:
await close_engine()
async def test_failed_run_reconciles_to_failed(self, tmp_path):
"""reconcile_stuck_once_tasks with terminal failed run → failed."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-recon-failed")
await _set_task_running(task_repo, "task-recon-failed")
await _create_run(run_repo, run_id="run-rf1", task_id="task-recon-failed", status="failed")
sf = get_session_factory()
assert sf is not None
async with sf() as session:
await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == "run-rf1").values(error="agent crashed"))
await session.commit()
count = await task_repo.reconcile_stuck_once_tasks(
error="interrupted: lease expired",
now=_NOW,
lease_grace_seconds=10,
)
assert count == 1
task = await _get_task(task_repo, "task-recon-failed")
assert task is not None
assert task["status"] == "failed"
assert task["last_error"] == "agent crashed"
finally:
await close_engine()
async def test_no_terminal_run_gets_generic_cancel(self, tmp_path):
"""reconcile_stuck_once_tasks with no terminal run → cancelled."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-recon-noterm")
await _set_task_running(task_repo, "task-recon-noterm")
# No run row at all.
count = await task_repo.reconcile_stuck_once_tasks(
error="interrupted: lease expired",
now=_NOW,
lease_grace_seconds=10,
)
assert count == 1
task = await _get_task(task_repo, "task-recon-noterm")
assert task is not None
assert task["status"] == "cancelled"
assert task["last_error"] == "interrupted: lease expired"
finally:
await close_engine()
async def test_reconcile_sees_run_committed_before_finalize(self, tmp_path):
"""Monkeypatch _fetch_latest_run to simulate the race: a completion
commits success after recovery has observed the stale 'running' state.
The test proves the post-lock fresh read picks up the committed success.
A reverted pre-lock batch implementation never calls _fetch_latest_run,
so the monkeypatch never fires and the test fails catching the regression.
"""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-recon-race")
await _set_task_running(task_repo, "task-recon-race")
await _create_run(run_repo, run_id="run-rrace", task_id="task-recon-race", status="running")
original_fetch = ScheduledTaskRepository._fetch_latest_run
fresh_sf = get_session_factory()
assert fresh_sf is not None
intercepted = False
async def _intercepted_fetch(session, task_id: str):
nonlocal intercepted
if not intercepted and task_id == "task-recon-race":
intercepted = True
async with fresh_sf() as cs:
await cs.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == "run-rrace").values(status="success"))
await cs.commit()
return await original_fetch(session, task_id)
return await original_fetch(session, task_id)
ScheduledTaskRepository._fetch_latest_run = staticmethod(_intercepted_fetch)
try:
count = await task_repo.reconcile_stuck_once_tasks(
error="interrupted: lease expired",
now=_NOW,
lease_grace_seconds=10,
)
assert count == 1
task = await _get_task(task_repo, "task-recon-race")
assert task is not None
assert task["status"] == "completed"
assert task["last_error"] is None
finally:
ScheduledTaskRepository._fetch_latest_run = staticmethod(original_fetch)
finally:
await close_engine()
# ---------------------------------------------------------------------------
# Regression: multiple historical runs for both reconciliation paths
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("recovery_method", ["cancel_stuck_once_tasks", "reconcile_stuck_once_tasks"])
@pytest.mark.parametrize("newer_status", ["skipped", "queued", "launching", "running"])
@pytest.mark.parametrize("older_clock_ahead_seconds", [30, 0], ids=["reversed-timestamps", "equal-timestamps"])
async def test_once_recovery_uses_occurrence_order_despite_clock_skew(tmp_path, recovery_method, newer_status, older_clock_ahead_seconds):
"""A newer admitted occurrence wins even when timestamps and IDs favor the older one."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo)
await _set_task_running(task_repo, "task-once-1")
# Exercise normal repository insertion; only the worker's clock changes.
# Descending UUID order must not break ties in favor of the old success.
with patch("deerflow.persistence.scheduled_task_runs.sql.datetime") as clock:
clock.now.return_value = _NOW + timedelta(seconds=older_clock_ahead_seconds)
older = await run_repo.create(
run_record_id="ffffffff-ffff-4fff-8fff-ffffffffffff",
task_id="task-once-1",
thread_id="thread-old",
scheduled_for=clock.now.return_value,
trigger="manual",
status="success",
)
clock.now.return_value = _NOW
newer = await run_repo.create(
run_record_id="00000000-0000-4000-8000-000000000000",
task_id="task-once-1",
thread_id="thread-new",
scheduled_for=clock.now.return_value,
trigger="manual",
status=newer_status,
)
assert older["created_at"] >= newer["created_at"]
assert older["scheduled_for"] >= newer["scheduled_for"]
kwargs = {"error": "interrupted: recovery"}
if recovery_method == "reconcile_stuck_once_tasks":
kwargs["now"] = _NOW + timedelta(minutes=1)
count = await getattr(task_repo, recovery_method)(**kwargs)
task = await _get_task(task_repo, "task-once-1")
assert task is not None
assert task["status"] == ("cancelled" if newer_status == "skipped" else "running")
assert count == (1 if newer_status == "skipped" else 0)
assert task["last_error"] is None
async with run_repo._sf() as session:
older_row = await session.get(ScheduledTaskRunRow, older["id"])
newer_row = await session.get(ScheduledTaskRunRow, newer["id"])
assert newer_row.occurrence_seq > older_row.occurrence_seq
finally:
await close_engine()
class TestCancelStuckMultipleRuns:
"""Regression: cancel_stuck_once_tasks with multiple historical runs.
Ensures the latest run row is used for finalisation, not an older one.
"""
async def test_older_success_newer_skipped_marks_cancelled(self, tmp_path):
"""Older success + newer skipped → task should be cancelled (not completed)."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-multi-cs")
await _set_task_running(task_repo, "task-multi-cs")
# Older run succeeded.
await _create_run(
run_repo,
run_id="run-old",
task_id="task-multi-cs",
status="success",
created_at=_NOW - timedelta(minutes=5),
)
# Newer run was skipped — the latest run determines the outcome.
await _create_run(
run_repo,
run_id="run-new",
task_id="task-multi-cs",
status="skipped",
created_at=_NOW,
)
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 1
task = await _get_task(task_repo, "task-multi-cs")
assert task is not None
assert task["status"] == "cancelled"
finally:
await close_engine()
async def test_newer_created_at_with_earlier_scheduled_for_wins(self, tmp_path):
"""Newer created_at + earlier scheduled_for → latest run should still win."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-multi-cs-skew")
await _set_task_running(task_repo, "task-multi-cs-skew")
await _create_run(
run_repo,
run_id="run-old-scheduled-late",
task_id="task-multi-cs-skew",
status="success",
created_at=_NOW - timedelta(minutes=5),
scheduled_for=_NOW + timedelta(minutes=5),
)
await _create_run(
run_repo,
run_id="run-new-scheduled-early",
task_id="task-multi-cs-skew",
status="skipped",
created_at=_NOW,
scheduled_for=_NOW - timedelta(minutes=5),
)
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 1
task = await _get_task(task_repo, "task-multi-cs-skew")
assert task is not None
assert task["status"] == "cancelled"
finally:
await close_engine()
async def test_same_scheduled_for_uses_created_at_recency_tie_break(self, tmp_path):
"""Same scheduled_for + newer created_at skipped → task should be cancelled."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-multi-cs-tie")
await _set_task_running(task_repo, "task-multi-cs-tie")
same_scheduled_for = _NOW
await _create_run(
run_repo,
run_id="run-old",
task_id="task-multi-cs-tie",
status="success",
created_at=_NOW - timedelta(minutes=5),
scheduled_for=same_scheduled_for,
)
await _create_run(
run_repo,
run_id="run-new",
task_id="task-multi-cs-tie",
status="skipped",
created_at=_NOW,
scheduled_for=same_scheduled_for,
)
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 1
task = await _get_task(task_repo, "task-multi-cs-tie")
assert task is not None
assert task["status"] == "cancelled"
finally:
await close_engine()
async def test_older_success_newer_active_leaves_parent_unchanged(self, tmp_path):
"""Older success + newer running → latest run is active, parent left unchanged (not cancelled)."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-multi-ca")
await _set_task_running(task_repo, "task-multi-ca")
# Older run succeeded.
await _create_run(
run_repo,
run_id="run-old",
task_id="task-multi-ca",
status="success",
created_at=_NOW - timedelta(minutes=5),
)
# Newer run is still running — active occurrence left untouched.
await _create_run(
run_repo,
run_id="run-new",
task_id="task-multi-ca",
status="running",
created_at=_NOW,
)
count = await task_repo.cancel_stuck_once_tasks(error="interrupted: gateway restarted")
assert count == 0
task = await _get_task(task_repo, "task-multi-ca")
assert task is not None
assert task["status"] == "running"
finally:
await close_engine()
class TestReconcileStuckMultipleRuns:
"""Regression: reconcile_stuck_once_tasks with multiple historical runs.
Ensures the latest run row is used for finalisation, not an older one.
"""
async def test_same_scheduled_for_uses_created_at_recency_tie_break(self, tmp_path):
"""Same scheduled_for + newer created_at skipped → task should be cancelled."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-multi-rs-tie")
await _set_task_running(task_repo, "task-multi-rs-tie")
same_scheduled_for = _NOW
await _create_run(
run_repo,
run_id="run-old",
task_id="task-multi-rs-tie",
status="success",
created_at=_NOW - timedelta(minutes=5),
scheduled_for=same_scheduled_for,
)
await _create_run(
run_repo,
run_id="run-new",
task_id="task-multi-rs-tie",
status="skipped",
created_at=_NOW,
scheduled_for=same_scheduled_for,
)
count = await task_repo.reconcile_stuck_once_tasks(
error="interrupted: lease expired",
now=_NOW,
lease_grace_seconds=10,
)
assert count == 1
task = await _get_task(task_repo, "task-multi-rs-tie")
assert task is not None
assert task["status"] == "cancelled"
finally:
await close_engine()
async def test_newer_created_at_with_earlier_scheduled_for_wins(self, tmp_path):
"""Newer created_at + earlier scheduled_for → latest run should still win."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-multi-rs-skew")
await _set_task_running(task_repo, "task-multi-rs-skew")
await _create_run(
run_repo,
run_id="run-old-scheduled-late",
task_id="task-multi-rs-skew",
status="success",
created_at=_NOW - timedelta(minutes=5),
scheduled_for=_NOW + timedelta(minutes=5),
)
await _create_run(
run_repo,
run_id="run-new-scheduled-early",
task_id="task-multi-rs-skew",
status="skipped",
created_at=_NOW,
scheduled_for=_NOW - timedelta(minutes=5),
)
count = await task_repo.reconcile_stuck_once_tasks(
error="interrupted: lease expired",
now=_NOW,
lease_grace_seconds=10,
)
assert count == 1
task = await _get_task(task_repo, "task-multi-rs-skew")
assert task is not None
assert task["status"] == "cancelled"
finally:
await close_engine()
async def test_older_success_newer_skipped_marks_cancelled(self, tmp_path):
"""Older success + newer skipped → task should be cancelled (not completed)."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-multi-rs")
await _set_task_running(task_repo, "task-multi-rs")
# Older run succeeded.
await _create_run(
run_repo,
run_id="run-old",
task_id="task-multi-rs",
status="success",
created_at=_NOW - timedelta(minutes=5),
)
# Newer run was skipped.
await _create_run(
run_repo,
run_id="run-new",
task_id="task-multi-rs",
status="skipped",
created_at=_NOW,
)
count = await task_repo.reconcile_stuck_once_tasks(
error="interrupted: lease expired",
now=_NOW,
lease_grace_seconds=10,
)
assert count == 1
task = await _get_task(task_repo, "task-multi-rs")
assert task is not None
assert task["status"] == "cancelled"
finally:
await close_engine()
async def test_older_success_newer_active_leaves_parent_unchanged(self, tmp_path):
"""Older success + newer running → latest run is active, parent left unchanged (not cancelled)."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-multi-ra")
await _set_task_running(task_repo, "task-multi-ra")
# Older run succeeded.
await _create_run(
run_repo,
run_id="run-old",
task_id="task-multi-ra",
status="success",
created_at=_NOW - timedelta(minutes=5),
)
# Newer run is still running — not terminal, so parent left unchanged.
await _create_run(
run_repo,
run_id="run-new",
task_id="task-multi-ra",
status="running",
created_at=_NOW,
)
count = await task_repo.reconcile_stuck_once_tasks(
error="interrupted: lease expired",
now=_NOW,
lease_grace_seconds=10,
)
assert count == 0
task = await _get_task(task_repo, "task-multi-ra")
assert task is not None
assert task["status"] == "running"
finally:
await close_engine()
# ---------------------------------------------------------------------------
# handle_run_completion happy path (Bug 1 guard)
# ---------------------------------------------------------------------------
class TestHandleRunCompletionHappyPath:
"""Ensure handle_run_completion correctly finalises a once task on success."""
async def test_success_completes_once_task(self, tmp_path):
"""Call the real handle_run_completion and verify the once task transitions to completed."""
task_repo, run_repo = await _init_db(tmp_path)
try:
await _create_once_task(task_repo, task_id="task-completion")
await _set_task_running(task_repo, "task-completion")
# The run row is in 'running' — simulating the state just before
# the completion hook fires.
await _create_run(run_repo, run_id="run-comp", task_id="task-completion", status="running")
service = ScheduledTaskService(
task_repo=task_repo,
task_run_repo=run_repo,
launch_run=lambda **_kw: None,
poll_interval_seconds=5,
lease_seconds=120,
max_concurrent_runs=3,
)
record = RunRecord(
run_id="run-comp",
thread_id="thread-1",
assistant_id=None,
status=RunStatus.success,
on_disconnect=DisconnectMode.continue_,
metadata={
"scheduled_task_id": "task-completion",
"scheduled_task_run_id": "run-comp",
},
user_id="user-1",
)
await service.handle_run_completion(record)
task = await _get_task(task_repo, "task-completion")
assert task is not None
assert task["status"] == "completed"
assert task["last_error"] is None
finally:
await close_engine()

View File

@ -0,0 +1,348 @@
"""Late launches must preserve the latest occurrence's parent projection.
The scheduler and repositories are real. Only the external run lifecycle and
barriers between its launch return and scheduler bookkeeping are controlled.
"""
from __future__ import annotations
import asyncio
import os
import uuid
from contextlib import suppress
from datetime import UTC, datetime, timedelta
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import pytest
import pytest_asyncio
from sqlalchemy import event, select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.schema import CreateSchema, DropSchema
from app.scheduler.service import ScheduledTaskService
from deerflow.persistence.base import Base
from deerflow.persistence.run import RunRepository
from deerflow.persistence.run.model import RunRow
from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow
from deerflow.runtime.runs.manager import RunRecord
from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus
pytestmark = pytest.mark.asyncio
def _asyncpg_url(url: str | None) -> str | None:
"""CI passes a sync ``postgresql://...?sslmode=disable`` URL; use asyncpg and drop libpq-only keys."""
if not url:
return url
parts = urlsplit(url)
scheme = "postgresql+asyncpg" if parts.scheme in {"postgres", "postgresql"} else parts.scheme
query = urlencode([(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key not in {"sslmode", "channel_binding"}])
return urlunsplit(parts._replace(scheme=scheme, query=query))
@pytest_asyncio.fixture(params=["sqlite", "postgres-single", "postgres-multi"])
async def occurrence_databases(request, tmp_path):
backend = request.param
postgres_uri = _asyncpg_url(os.environ.get("TEST_POSTGRES_URI"))
if backend != "sqlite" and not postgres_uri:
pytest.skip("TEST_POSTGRES_URI is not set")
schema = "scheduler_order_" + uuid.uuid4().hex
admin = None
engines = []
try:
if backend != "sqlite":
admin = create_async_engine(postgres_uri)
async with admin.begin() as connection:
await connection.execute(CreateSchema(schema))
for _ in range(3):
if backend == "sqlite":
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'scheduler.db'}")
@event.listens_for(engine.sync_engine, "connect")
def configure_sqlite(connection, _record):
cursor = connection.cursor()
cursor.execute("PRAGMA journal_mode=WAL")
cursor.execute("PRAGMA busy_timeout=10000")
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()
else:
engine = create_async_engine(postgres_uri, connect_args={"server_settings": {"search_path": schema}})
engines.append(engine)
async with engines[0].begin() as connection:
await connection.run_sync(lambda sync: Base.metadata.create_all(sync, tables=[ScheduledTaskRow.__table__, ScheduledTaskRunRow.__table__, RunRow.__table__]))
yield [async_sessionmaker(engine, expire_on_commit=False) for engine in engines], backend == "postgres-multi"
finally:
for engine in engines:
await engine.dispose()
if admin is not None:
async with admin.begin() as connection:
await connection.execute(DropSchema(schema, cascade=True, if_exists=True))
await admin.dispose()
class _RevivalBarrierRepository(ScheduledTaskRunRepository):
def __init__(self, factory, durable, revived, resume):
super().__init__(factory, run_repository=durable)
self.revived = revived
self.resume = resume
async def reconcile_launched_run(self, *args, **kwargs):
result = await super().reconcile_launched_run(*args, **kwargs)
assert result is True
self.revived.set()
await self.resume.wait()
return result
def _service(tasks, occurrences, launcher, multi):
return ScheduledTaskService(
task_repo=tasks,
task_run_repo=occurrences,
launch_run=launcher,
poll_interval_seconds=3600,
lease_seconds=5,
max_concurrent_runs=3,
queue_timeout_seconds=10,
multi_instance=multi,
run_lease_grace_seconds=0,
)
async def _parent_projection(factory, task_id):
async with factory() as session:
parent = await session.get(ScheduledTaskRow, task_id)
return {key: getattr(parent, key) for key in ("status", "last_error", "last_run_id", "last_thread_id", "last_run_at", "next_run_at", "lease_owner", "lease_expires_at", "run_count")}
def _completion(run_id, launch_args, status, error=None):
return RunRecord(
run_id=run_id,
thread_id=launch_args["thread_id"],
assistant_id=None,
status=status,
on_disconnect=DisconnectMode.continue_,
user_id="user-order",
metadata=launch_args["metadata"],
error=error,
)
@pytest.mark.parametrize("interrupt_before_parent", [True, False], ids=["restart-after-revival", "late-completion-callback"])
@pytest.mark.parametrize("newer_occurrence", [True, False], ids=["newer-failed-occurrence", "single-occurrence-control"])
async def test_late_launch_preserves_newer_result_and_counts_once(occurrence_databases, interrupt_before_parent, newer_occurrence):
factories, multi = occurrence_databases
durable_a, durable_b, durable_recovery = [RunRepository(factory) for factory in factories]
tasks_a, tasks_b, tasks_recovery = [ScheduledTaskRepository(factory, run_repository=durable) for factory, durable in zip(factories, (durable_a, durable_b, durable_recovery), strict=True)]
entered, release, revived, resume = [asyncio.Event() for _ in range(4)]
occurrences_a = _RevivalBarrierRepository(factories[0], durable_a, revived, resume)
occurrences_b = ScheduledTaskRunRepository(factories[1], run_repository=durable_b)
occurrences_recovery = ScheduledTaskRunRepository(factories[2], run_repository=durable_recovery)
now = datetime.now(UTC)
task_id, run_a, run_b = "task-order", "run-order-a", "run-order-b"
launches = {}
async def launch_a(**kwargs):
launches["a"] = kwargs
entered.set()
await release.wait()
# Model a real run becoming terminal before its slow launch call
# returns. Scheduled occurrence states are only changed by the service.
await durable_a.put(run_a, thread_id=kwargs["thread_id"], user_id="user-order", status="running", metadata=kwargs["metadata"])
await durable_a.update_status(run_a, "success")
return {"run_id": run_a, "thread_id": kwargs["thread_id"]}
async def launch_b(**kwargs):
launches["b"] = kwargs
await durable_b.put(run_b, thread_id=kwargs["thread_id"], user_id="user-order", status="running", metadata=kwargs["metadata"])
return {"run_id": run_b, "thread_id": kwargs["thread_id"]}
async def singleton_launcher(**kwargs):
if "a" not in launches:
return await launch_a(**kwargs)
assert "b" not in launches
return await launch_b(**kwargs)
async def unexpected_launch(**_kwargs):
pytest.fail("recovery must not launch another occurrence")
service_a = _service(tasks_a, occurrences_a, launch_a if multi else singleton_launcher, multi)
service_b = _service(tasks_b, occurrences_b, launch_b, multi) if multi else service_a
peer = _service(tasks_recovery, occurrences_recovery, unexpected_launch, multi) if multi else service_a
scheduled_at = now - timedelta(seconds=1) if multi else now + timedelta(days=1)
await tasks_a.create(
task_id=task_id,
user_id="user-order",
thread_id=None,
context_mode="fresh_thread_per_run",
assistant_id=None,
title="occurrence order",
prompt="test",
schedule_type="once",
schedule_spec={"run_at": scheduled_at.isoformat()},
timezone="UTC",
next_run_at=scheduled_at,
)
initial_task = await tasks_a.get(task_id, user_id="user-order")
dispatch = asyncio.create_task(service_a.run_once(now=now) if multi else service_a.dispatch_task(initial_task, now=now, trigger="manual"))
restarted = _service(tasks_recovery, occurrences_recovery, unexpected_launch, multi)
async def parked_poll_loop():
await restarted._stop.wait()
restarted._run_loop = parked_poll_loop
try:
await asyncio.wait_for(entered.wait(), 10)
assert (await occurrences_b.list_by_task(task_id))[0]["status"] == "launching"
# Explicit scheduler time expires the launch claim and then the queue;
# no sleeps or direct terminal occurrence writes are needed.
await peer.run_once(now=now + timedelta(seconds=30))
expired = (await occurrences_b.list_by_task(task_id))[0]
assert expired["status"] == "failed"
assert expired["run_id"] is None
if newer_occurrence:
result = await service_b.dispatch_task(await tasks_b.get(task_id, user_id="user-order"), now=now + timedelta(seconds=31), trigger="manual")
assert result["outcome"] == "launched"
await durable_b.update_status(run_b, "error", error="newer run failed")
await service_b.handle_run_completion(_completion(run_b, launches["b"], RunStatus.error, "newer run failed"))
expected = await _parent_projection(factories[1], task_id)
assert expected["status"] == "failed"
assert expected["last_run_id"] == run_b
assert expected["run_count"] == 1
expected["run_count"] = 2
release.set()
await asyncio.wait_for(revived.wait(), 10)
if interrupt_before_parent:
# The revival transaction committed, but dispatch has not yet
# performed the parent launch accounting.
dispatch.cancel()
with suppress(asyncio.CancelledError):
await dispatch
else:
resume.set()
await asyncio.wait_for(dispatch, 10)
await service_a.handle_run_completion(_completion(run_a, launches["a"], RunStatus.success))
if newer_occurrence:
assert await _parent_projection(factories[1], task_id) == expected
await restarted.start()
await restarted.stop()
after = await _parent_projection(factories[1], task_id)
if newer_occurrence:
assert after == expected
else:
assert after["status"] == "completed"
assert after["last_run_id"] == run_a
assert after["run_count"] == 1
async with factories[1]() as session:
statuses = dict((await session.execute(select(ScheduledTaskRunRow.run_id, ScheduledTaskRunRow.status).where(ScheduledTaskRunRow.task_id == task_id))).all())
assert statuses == ({run_a: "success", run_b: "failed"} if newer_occurrence else {run_a: "success"})
# Repeated completion and recovery must neither re-project old state
# nor account either successful launch a second time.
await service_a.handle_run_completion(_completion(run_a, launches["a"], RunStatus.success))
assert await _parent_projection(factories[1], task_id) == after
if newer_occurrence:
await service_b.handle_run_completion(_completion(run_b, launches["b"], RunStatus.error, "newer run failed"))
assert await _parent_projection(factories[1], task_id) == after
await restarted.start()
await restarted.stop()
assert await _parent_projection(factories[1], task_id) == after
finally:
release.set()
resume.set()
if not dispatch.done():
dispatch.cancel()
with suppress(asyncio.CancelledError):
await dispatch
await restarted.stop()
@pytest.mark.parametrize("interrupt_before_return", [True, False], ids=["crash-before-launch-return", "normal-launch-return"])
async def test_completion_before_launch_return_accounts_once(occurrence_databases, interrupt_before_return):
factories, multi = occurrence_databases
durable = RunRepository(factories[0])
tasks = ScheduledTaskRepository(factories[0], run_repository=durable)
occurrences = ScheduledTaskRunRepository(factories[0], run_repository=durable)
completed, release = asyncio.Event(), asyncio.Event()
run_id, task_id = "run-fast-completion", "task-fast-completion"
now = datetime.now(UTC)
record = None
async def launcher(**kwargs):
nonlocal record
await durable.put(run_id, thread_id=kwargs["thread_id"], user_id="user-order", status="running", metadata=kwargs["metadata"])
await durable.update_status(run_id, "success")
record = _completion(run_id, kwargs, RunStatus.success)
await service.handle_run_completion(record)
completed.set()
await release.wait()
return {"run_id": run_id, "thread_id": kwargs["thread_id"]}
async def unexpected_launch(**_kwargs):
pytest.fail("completed occurrence must not launch again")
service = _service(tasks, occurrences, launcher, multi)
await tasks.create(
task_id=task_id,
user_id="user-order",
thread_id=None,
context_mode="fresh_thread_per_run",
assistant_id=None,
title="fast completion",
prompt="test",
schedule_type="once",
schedule_spec={"run_at": now.isoformat()},
timezone="UTC",
next_run_at=now,
)
dispatch = asyncio.create_task(service.dispatch_task(await tasks.get(task_id, user_id="user-order"), now=now, trigger="manual"))
recovery_durable = RunRepository(factories[1])
restarted = _service(
ScheduledTaskRepository(factories[1], run_repository=recovery_durable),
ScheduledTaskRunRepository(factories[1], run_repository=recovery_durable),
unexpected_launch,
multi,
)
async def parked_poll_loop():
await restarted._stop.wait()
restarted._run_loop = parked_poll_loop
try:
await asyncio.wait_for(completed.wait(), 10)
before_return = await _parent_projection(factories[1], task_id)
assert before_return["status"] == "completed"
assert before_return["run_count"] == 1
assert before_return["last_run_id"] == run_id
assert before_return["last_thread_id"] == record.thread_id
if interrupt_before_return:
dispatch.cancel()
with suppress(asyncio.CancelledError):
await dispatch
else:
release.set()
await asyncio.wait_for(dispatch, 10)
await restarted.start()
await restarted.stop()
after = await _parent_projection(factories[1], task_id)
assert after["status"] == "completed"
assert after["last_run_id"] == run_id
assert after["run_count"] == 1
assert (await occurrences.list_by_task(task_id))[0]["status"] == "success"
await service.handle_run_completion(record)
await restarted.start()
await restarted.stop()
assert await _parent_projection(factories[1], task_id) == after
finally:
release.set()
if not dispatch.done():
dispatch.cancel()
with suppress(asyncio.CancelledError):
await dispatch
await restarted.stop()