From ca3e510b7d55fb414c45db2679beb38c2eafa57b Mon Sep 17 00:00:00 2001 From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com> Date: Fri, 24 Jul 2026 06:41:09 -0700 Subject: [PATCH] fix(scheduler): close duplicate dispatch race (#4105) Enforce one queued or running scheduled-task run per task with a partial unique index. The migration resolves legacy duplicates before creating the index, and losing inserts use the existing conflict or skip outcomes. --- backend/AGENTS.md | 5 +- backend/app/scheduler/service.py | 96 ++++++-- .../0007_scheduled_run_active_index.py | 130 +++++++++++ .../scheduled_task_runs/__init__.py | 4 +- .../persistence/scheduled_task_runs/model.py | 31 ++- .../persistence/scheduled_task_runs/sql.py | 34 ++- ...est_migration_0004_run_ownership_dedupe.py | 3 +- ...ration_0007_scheduled_run_active_dedupe.py | 180 ++++++++++++++ backend/tests/test_persistence_bootstrap.py | 2 +- .../test_persistence_bootstrap_concurrency.py | 2 +- .../test_persistence_bootstrap_regression.py | 4 +- .../test_scheduled_task_dispatch_race.py | 219 ++++++++++++++++++ .../tests/test_scheduled_task_repository.py | 14 +- backend/tests/test_scheduled_task_service.py | 6 +- 14 files changed, 695 insertions(+), 35 deletions(-) create mode 100644 backend/packages/harness/deerflow/persistence/migrations/versions/0007_scheduled_run_active_index.py create mode 100644 backend/tests/test_migration_0007_scheduled_run_active_dedupe.py create mode 100644 backend/tests/test_scheduled_task_dispatch_race.py diff --git a/backend/AGENTS.md b/backend/AGENTS.md index e84534db8..424c4b16f 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -17,6 +17,7 @@ DeerFlow is a LangGraph-based AI super agent system with a full-stack architectu - Gateway streams `write_file` and `str_replace` argument deltas in bounded batches when clients also subscribe to `values`; messages-only consumers retain the original per-chunk contract, while `values` preserves the complete tool call. - With `stream_subgraphs`, subgraph frames keep their namespace in the SSE event name (`values|`, LangGraph Platform style) instead of impersonating root frames — a delegated subagent inherits the parent checkpoint namespace, so publishing its `values` snapshot as bare `values` replaces the whole thread view in SDK clients (#4399). Root-only consumers (file-tool chunk batcher, subagent event persistence, LLM error-fallback detection) ignore namespaced frames. The web frontend does not request subgraph streaming; subtask progress rides root-namespace `task_*` custom events. - Scheduled-task executions must reuse that same Gateway run lifecycle. The scheduler may decide *when* work runs, but it must dispatch through the existing run path rather than introducing a parallel execution stack. +- Scheduled-task dispatch enforces "at most one active run per task when `overlap_policy=skip`" at the DB layer via the partial unique index `uq_scheduled_task_run_active` (`scheduled_task_runs.task_id WHERE status IN ('queued','running')`). `ScheduledTaskService.dispatch_task`'s `has_active_runs` check is a non-atomic fast path (its own session, separated from the `create()` insert by `await` points), so two concurrent dispatches — a manual `POST /scheduled-tasks/{id}/trigger` racing the poller, a double-click, or a client retry — can both pass it; the index is the atomic arbiter, and the losing `create` surfaces as `ActiveScheduledRunConflict` (translated from `IntegrityError` in the repository) and collapses to the same outcome as the fast path (manual → 409 conflict, scheduled → a `"skipped"` tombstone). The scheduled-skip tombstone is created directly as terminal `"skipped"` (not a transient `"queued"`) so it never occupies the active slot the pre-existing run still holds. Sibling of the `runs` table's `uq_runs_thread_active` (PR #4003), which keys on `thread_id` and so does not cover the default `fresh_thread_per_run` context where every dispatch gets a new thread. Index is status-only, not `overlap_policy`-conditional (the policy is fixed to `"skip"` in the MVP). **Project Structure**: ``` @@ -810,8 +811,10 @@ This invokes `alembic revision --autogenerate` against the live ORM models. Revi - `migrations/_helpers.py` — `safe_add_column` / `safe_drop_column` - `migrations/versions/0001_baseline.py` — chain root, matches the schema `create_all` produces from `Base.metadata` - `migrations/versions/0002_runs_token_usage.py` — fixes issue #3682 +- `migrations/versions/0004_run_ownership.py` — `runs` multi-worker ownership + the `uq_runs_thread_active` partial unique index, with a `_dedupe_active_runs_per_thread()` pre-step so `CREATE UNIQUE INDEX` cannot fail on a field DB that already has duplicate active rows per thread +- `migrations/versions/0007_scheduled_run_active_index.py` — the `uq_scheduled_task_run_active` partial unique index (at most one queued/running `scheduled_task_runs` row per `task_id`), with a `_dedupe_active_scheduled_runs_per_task()` pre-step (keeps the newest active row per task, supersedes the rest to `interrupted` with an explanatory `error` + `finished_at`) mirroring 0004; chains after `0006_agents` - `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch decision + locking -- 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), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor) +- 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), `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) ### Checkpoint Channel Modes (`full` / `delta`) diff --git a/backend/app/scheduler/service.py b/backend/app/scheduler/service.py index b294c8957..3080fe008 100644 --- a/backend/app/scheduler/service.py +++ b/backend/app/scheduler/service.py @@ -9,11 +9,17 @@ from typing import Any, Literal from fastapi import HTTPException +from deerflow.persistence.scheduled_task_runs import ActiveScheduledRunConflict from deerflow.runtime import ConflictError, RunRecord from deerflow.scheduler.schedules import next_run_at logger = logging.getLogger(__name__) +# Shared so the has_active_runs fast path and the unique-index race path return +# byte-identical outcomes for the same "task already has an active run" condition. +_ACTIVE_RUN_CONFLICT_ERROR = "task already has an active run" +_SKIP_ACTIVE_RUN_ERROR = "skipped: a previous run of this task is still active" + class ScheduledTaskService: def __init__( @@ -94,28 +100,37 @@ class ScheduledTaskService: # does not count itself as the active run. A manual trigger against an # active run is rejected outright (409 at the router) instead of being # recorded as a skipped occurrence — nothing was scheduled to happen. - skip_error: str | None = None - if task.get("overlap_policy", "skip") == "skip" and await self._task_run_repo.has_active_runs(task["id"]): + # + # This has_active_runs check is a non-atomic fast path: it runs in its + # own session and is separated from the create() below by await points, + # so two concurrent dispatches (double-click / client retry / a manual + # trigger racing the poller) can both observe no active run. The DB is + # the atomic arbiter — the partial unique index uq_scheduled_task_run_active + # rejects the second active insert, surfaced as ActiveScheduledRunConflict + # and collapsed to the SAME outcome as this fast path just below. + overlap_skip = task.get("overlap_policy", "skip") == "skip" + if overlap_skip and await self._task_run_repo.has_active_runs(task["id"]): if trigger == "manual": - return { - "outcome": "conflict", - "task_run_id": None, - "run_id": None, - "thread_id": execution_thread_id, - "error": "task already has an active run", - } - skip_error = "skipped: a previous run of this task is still active" + return self._active_run_conflict_result(execution_thread_id) + return await self._record_scheduled_skip(task, thread_id=execution_thread_id, now=now, trigger=trigger) + task_run_id = f"task-run-{uuid.uuid4().hex}" - await self._task_run_repo.create( - run_record_id=task_run_id, - task_id=task["id"], - thread_id=execution_thread_id, - scheduled_for=now, - trigger=trigger, - status="queued", - ) - if skip_error is not None: - return await self._finalize_skip(task, task_run_id=task_run_id, thread_id=execution_thread_id, now=now, error=skip_error) + try: + await self._task_run_repo.create( + run_record_id=task_run_id, + task_id=task["id"], + thread_id=execution_thread_id, + scheduled_for=now, + trigger=trigger, + status="queued", + ) + except ActiveScheduledRunConflict: + # Lost the create race for the task's single active slot: a + # concurrent dispatch passed the same fast-path check and inserted + # its active row first. Identical outcome to the fast path above. + if trigger == "manual": + return self._active_run_conflict_result(execution_thread_id) + return await self._record_scheduled_skip(task, thread_id=execution_thread_id, now=now, trigger=trigger) try: result = await self._launch_run( thread_id=execution_thread_id, @@ -209,6 +224,47 @@ class ScheduledTaskService: "error": str(exc), } + def _active_run_conflict_result(self, thread_id: str) -> dict[str, Any]: + """Manual-trigger response when the task already has an active run. + + Nothing was scheduled to happen, so no run-history row is recorded; the + router maps this to a 409. + """ + return { + "outcome": "conflict", + "task_run_id": None, + "run_id": None, + "thread_id": thread_id, + "error": _ACTIVE_RUN_CONFLICT_ERROR, + } + + async def _record_scheduled_skip( + self, + task: dict[str, Any], + *, + thread_id: str, + now: datetime, + trigger: str, + ) -> dict[str, Any]: + """Record a skipped occurrence for a scheduled dispatch that overlapped an active run. + + The tombstone is created directly as terminal ``"skipped"`` rather than + the transient ``"queued"`` the launch path uses: a queued row is active + and would itself trip ``uq_scheduled_task_run_active`` against the + pre-existing run that is still holding the task's single active slot. + ``"skipped"`` is outside the index predicate, so it never conflicts. + """ + task_run_id = f"task-run-{uuid.uuid4().hex}" + await self._task_run_repo.create( + run_record_id=task_run_id, + task_id=task["id"], + thread_id=thread_id, + scheduled_for=now, + trigger=trigger, + status="skipped", + ) + return await self._finalize_skip(task, task_run_id=task_run_id, thread_id=thread_id, now=now, error=_SKIP_ACTIVE_RUN_ERROR) + async def _finalize_skip( self, task: dict[str, Any], diff --git a/backend/packages/harness/deerflow/persistence/migrations/versions/0007_scheduled_run_active_index.py b/backend/packages/harness/deerflow/persistence/migrations/versions/0007_scheduled_run_active_index.py new file mode 100644 index 000000000..120a076b1 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/versions/0007_scheduled_run_active_index.py @@ -0,0 +1,130 @@ +"""scheduled task run active uniqueness. + +Revision ID: 0007_scheduled_run_active_index +Revises: 0006_agents +Create Date: 2026-07-11 +""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from datetime import UTC, datetime + +import sqlalchemy as sa +from alembic import op + +logger = logging.getLogger(__name__) + +revision: str = "0007_scheduled_run_active_index" +down_revision: str | Sequence[str] | None = "0006_agents" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _dedupe_active_scheduled_runs_per_task() -> None: + """Supersede duplicate active rows so the partial unique index can be built. + + ``uq_scheduled_task_run_active`` enforces at most one queued/running row per + ``task_id``. A DB that already has two+ active rows for the same task (the + exact TOCTOU this PR closes: two concurrent ``dispatch_task`` calls both + passed ``has_active_runs`` and both inserted a "queued" row) would fail + ``CREATE UNIQUE INDEX`` and abort the alembic upgrade, blocking gateway + startup. + + Keep the newest active row per ``task_id`` (by ``created_at`` DESC, ``id`` + DESC as a deterministic tiebreaker) and mark the rest ``interrupted`` with + an explanatory ``error`` and a ``finished_at`` — the same orphan semantics + ``ScheduledTaskRunRepository.mark_stale_active_runs`` uses for runs whose + process is gone. + """ + bind = op.get_bind() + superseded_message = "interrupted during migration 0007_scheduled_run_active_index: superseded by a newer active run for the same scheduled task (partial unique index uq_scheduled_task_run_active)" + find_dupe_rows = sa.text( + """ + SELECT id, task_id + FROM scheduled_task_runs AS r1 + WHERE r1.status IN ('queued', 'running') + AND EXISTS ( + SELECT 1 FROM scheduled_task_runs AS r2 + WHERE r2.task_id = r1.task_id + AND r2.status IN ('queued', 'running') + AND r2.id <> r1.id + AND ( + r2.created_at > r1.created_at + OR (r2.created_at = r1.created_at AND r2.id > r1.id) + ) + ) + """ + ) + rows = list(bind.execute(find_dupe_rows).fetchall()) + if not rows: + return + for run_id, task_id in rows: + logger.warning( + "migration 0007_scheduled_run_active_index: superseding duplicate active scheduled run %s on task %s", + run_id, + task_id, + ) + update_stmt = sa.text( + """ + UPDATE scheduled_task_runs + SET status = 'interrupted', + error = :error_message, + finished_at = :finished_at + WHERE status IN ('queued', 'running') + AND EXISTS ( + SELECT 1 FROM scheduled_task_runs AS r2 + WHERE r2.task_id = scheduled_task_runs.task_id + AND r2.status IN ('queued', 'running') + AND r2.id <> scheduled_task_runs.id + AND ( + r2.created_at > scheduled_task_runs.created_at + OR (r2.created_at = scheduled_task_runs.created_at AND r2.id > scheduled_task_runs.id) + ) + ) + """ + ).bindparams( + sa.bindparam("error_message"), + # Typed so SQLAlchemy applies the dialect's DateTime bind processor + # (SQLite string format / Postgres timestamptz) instead of handing a + # raw datetime to the DBAPI (Python 3.12 dropped sqlite3's default + # datetime adapter). + sa.bindparam("finished_at", type_=sa.DateTime(timezone=True)), + ) + bind.execute( + update_stmt, + {"error_message": superseded_message, "finished_at": datetime.now(UTC)}, + ) + + +def upgrade() -> None: + # Idempotent index creation: the legacy/empty bootstrap path runs + # create_all (which creates the index from the ORM __table_args__) before + # upgrade head, so the migration must not fail when the index already + # exists. + insp = sa.inspect(op.get_bind()) + existing = {ix["name"] for ix in insp.get_indexes("scheduled_task_runs")} + if "uq_scheduled_task_run_active" not in existing: + # Supersede duplicate active rows first so the partial UNIQUE index can + # be built on DBs that already violate the invariant. No-op on clean + # DBs (the common path -- create_all already created the index, so this + # branch only runs on legacy DBs that pre-date the index). + _dedupe_active_scheduled_runs_per_task() + with op.batch_alter_table("scheduled_task_runs", schema=None) as batch_op: + batch_op.create_index( + "uq_scheduled_task_run_active", + ["task_id"], + unique=True, + sqlite_where=sa.text("status IN ('queued', 'running')"), + postgresql_where=sa.text("status IN ('queued', 'running')"), + ) + + +def downgrade() -> None: + bind = op.get_bind() + insp = sa.inspect(bind) + existing = {ix["name"] for ix in insp.get_indexes("scheduled_task_runs")} + if "uq_scheduled_task_run_active" in existing: + with op.batch_alter_table("scheduled_task_runs", schema=None) as batch_op: + batch_op.drop_index("uq_scheduled_task_run_active") diff --git a/backend/packages/harness/deerflow/persistence/scheduled_task_runs/__init__.py b/backend/packages/harness/deerflow/persistence/scheduled_task_runs/__init__.py index 3f9edfbf0..8c014cf90 100644 --- a/backend/packages/harness/deerflow/persistence/scheduled_task_runs/__init__.py +++ b/backend/packages/harness/deerflow/persistence/scheduled_task_runs/__init__.py @@ -1,4 +1,4 @@ from .model import ScheduledTaskRunRow -from .sql import ScheduledTaskRunRepository +from .sql import ActiveScheduledRunConflict, ScheduledTaskRunRepository -__all__ = ["ScheduledTaskRunRow", "ScheduledTaskRunRepository"] +__all__ = ["ActiveScheduledRunConflict", "ScheduledTaskRunRow", "ScheduledTaskRunRepository"] diff --git a/backend/packages/harness/deerflow/persistence/scheduled_task_runs/model.py b/backend/packages/harness/deerflow/persistence/scheduled_task_runs/model.py index dd165521e..cd352e823 100644 --- a/backend/packages/harness/deerflow/persistence/scheduled_task_runs/model.py +++ b/backend/packages/harness/deerflow/persistence/scheduled_task_runs/model.py @@ -2,7 +2,7 @@ from __future__ import annotations from datetime import UTC, datetime -from sqlalchemy import DateTime, String, Text +from sqlalchemy import DateTime, Index, String, Text, text from sqlalchemy.orm import Mapped, mapped_column from deerflow.persistence.base import Base @@ -22,3 +22,32 @@ class ScheduledTaskRunRow(Base): started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) + + __table_args__ = ( + # At most one active (queued/running) run per task. This is the atomic + # arbiter for the ``dispatch_task`` skip policy: the non-atomic + # ``has_active_runs`` check-then-create is a fast path, but two + # concurrent dispatches (double-click / client retry / manual trigger + # racing the poller) can both pass it, so the DB must reject the second + # active insert. Sibling of the ``runs`` table's ``uq_runs_thread_active`` + # (PR #4003); that one keys on ``thread_id`` and does not cover the + # default ``fresh_thread_per_run`` context (every dispatch gets a new + # thread), which is why the scheduled-task run row needs its own guard. + # + # Condition is status-only, not ``overlap_policy``: the policy is fixed + # to "skip" in the MVP, so a status-only predicate enforces the current + # invariant without denormalizing ``overlap_policy`` onto the run row + # for an unimplemented non-skip policy. If a non-skip policy is added + # this must become conditional (e.g. ``... AND overlap_policy = 'skip'``). + # + # Must live in ORM ``__table_args__`` (not just the migration) because + # the empty-DB bootstrap path runs ``create_all`` + ``stamp head`` and + # never executes the migration that also defines this index. + Index( + "uq_scheduled_task_run_active", + "task_id", + unique=True, + sqlite_where=text("status IN ('queued', 'running')"), + postgresql_where=text("status IN ('queued', 'running')"), + ), + ) diff --git a/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py b/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py index 4f32e90c2..afedb4b4f 100644 --- a/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py +++ b/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py @@ -4,6 +4,7 @@ from datetime import UTC, datetime from typing import Any from sqlalchemy import func, select +from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow @@ -13,6 +14,26 @@ TERMINAL_RUN_STATUSES: frozenset[str] = frozenset({"success", "failed", "skipped ACTIVE_RUN_STATUSES: tuple[str, ...] = ("queued", "running") +class ActiveScheduledRunConflict(Exception): + """A concurrent dispatch already holds the task's single active-run slot. + + Raised by :meth:`ScheduledTaskRunRepository.create` when inserting an + active (queued/running) run row would violate the partial unique index + ``uq_scheduled_task_run_active`` (at most one active run per ``task_id``). + This is the atomic counterpart to the non-atomic ``has_active_runs`` check + in ``ScheduledTaskService.dispatch_task``: two dispatches can both pass that + check, but only one can insert the active row — the loser lands here. + + Translating the SQLAlchemy ``IntegrityError`` into a domain exception at + the repository boundary keeps the service layer free of ``sqlalchemy.exc`` + coupling (mirrors ``deerflow.runtime.ConflictError`` for the runs table). + """ + + def __init__(self, task_id: str) -> None: + self.task_id = task_id + super().__init__(f"scheduled task {task_id!r} already has an active run") + + class ScheduledTaskRunRepository: def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None: self._sf = session_factory @@ -46,7 +67,18 @@ class ScheduledTaskRunRepository: ) async with self._sf() as session: session.add(row) - await session.commit() + try: + await session.commit() + except IntegrityError: + await session.rollback() + # Only active-status inserts can trip the partial unique index + # ``uq_scheduled_task_run_active``; a terminal-status row (e.g. + # a "skipped" tombstone) is outside its predicate and cannot + # conflict, so any IntegrityError there is a genuine fault and + # is re-raised untranslated. + if status in ACTIVE_RUN_STATUSES: + raise ActiveScheduledRunConflict(task_id) from None + raise await session.refresh(row) return self._row_to_dict(row) diff --git a/backend/tests/test_migration_0004_run_ownership_dedupe.py b/backend/tests/test_migration_0004_run_ownership_dedupe.py index 28bbd611a..25a75f539 100644 --- a/backend/tests/test_migration_0004_run_ownership_dedupe.py +++ b/backend/tests/test_migration_0004_run_ownership_dedupe.py @@ -156,7 +156,8 @@ async def test_migration_dedupes_duplicate_active_rows_before_unique_index(tmp_p with sqlite3.connect(db_path) as raw: version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0006_agents" + # Bootstrap upgrades through the later revisions after 0004. + assert version_row[0] == "0007_scheduled_run_active_index" # Sanity: the invariant the index enforces is now true — at most one # active row per thread. diff --git a/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py b/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py new file mode 100644 index 000000000..e737697c2 --- /dev/null +++ b/backend/tests/test_migration_0007_scheduled_run_active_dedupe.py @@ -0,0 +1,180 @@ +"""Regression test for migration ``0007_scheduled_run_active_index`` dedupe pass. + +End-to-end shape (mirrors ``test_migration_0004_run_ownership_dedupe``): + +1. Hand-build a SQLite DB that mirrors a real pre-0007 deployment that ran the + racy ``dispatch_task`` check-then-insert and accumulated duplicate active + rows per ``task_id`` (the exact dirty state the partial unique index + targets). +2. Stamp it at ``0006_agents`` so ``bootstrap_schema`` takes the + versioned branch and runs ``alembic upgrade head``. +3. Insert two+ queued/running rows for the same ``task_id`` (only possible + because the partial unique index does not exist yet). +4. Run ``init_engine`` (the FastAPI lifespan entry point), which routes through + ``bootstrap_schema`` -> ``upgrade head`` -> ``0007.upgrade()``. +5. Verify the migration superseded the older duplicates (set them to + ``interrupted`` with an explanatory message + ``finished_at``), kept the + newest active row, and successfully built the ``uq_scheduled_task_run_active`` + partial unique index. + +Pre-fix codepath would have raised ``UNIQUE constraint failed`` (SQLite) / +``could not create unique index`` (Postgres) on step 5, aborting the alembic +upgrade and blocking gateway startup. +""" + +from __future__ import annotations + +import sqlite3 +from datetime import UTC, datetime, timedelta +from pathlib import Path + +import pytest +import sqlalchemy as sa +from sqlalchemy.orm import Session + +import deerflow.persistence.models # noqa: F401 -- registers ORM models +from deerflow.persistence.base import Base +from deerflow.persistence.engine import close_engine, init_engine +from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow + +pytestmark = pytest.mark.asyncio + + +def _seed_pre_0007_with_duplicates(db_path: Path) -> None: + """Build a DB at revision 0006 with duplicate active rows per task_id. + + ``Base.metadata.create_all`` produces the full current schema (including + the partial unique index), so we drop just that index to land in the dirty + state the migration's dedupe pass targets, then stamp at 0005 and insert + the duplicates via the ORM so Python-side defaults populate. + """ + db_path.parent.mkdir(parents=True, exist_ok=True) + sync_engine = sa.create_engine(f"sqlite:///{db_path.as_posix()}") + try: + Base.metadata.create_all(sync_engine) + with sync_engine.begin() as conn: + # Drop only the partial unique index — its absence is what permits + # duplicate active rows to exist in the first place. + conn.execute(sa.text("DROP INDEX IF EXISTS uq_scheduled_task_run_active")) + # Stamp at 0006 so bootstrap takes the versioned branch and runs + # ``alembic upgrade head`` (which is what executes 0007.upgrade()). + conn.execute(sa.text("CREATE TABLE IF NOT EXISTS alembic_version (version_num VARCHAR(32) NOT NULL)")) + conn.execute(sa.text("DELETE FROM alembic_version")) + conn.execute(sa.text("INSERT INTO alembic_version (version_num) VALUES ('0006_agents')")) + + base = datetime.now(UTC) + with Session(sync_engine) as session: + session.add_all( + [ + # Task with three active rows: two older duplicates + newest. + ScheduledTaskRunRow( + id="run-old-a", + task_id="task-dup", + thread_id="thread-a", + scheduled_for=base, + trigger="scheduled", + status="queued", + created_at=base, + ), + ScheduledTaskRunRow( + id="run-old-b", + task_id="task-dup", + thread_id="thread-b", + scheduled_for=base, + trigger="manual", + status="running", + created_at=base + timedelta(seconds=10), + ), + ScheduledTaskRunRow( + id="run-newest", + task_id="task-dup", + thread_id="thread-c", + scheduled_for=base, + trigger="scheduled", + status="queued", + created_at=base + timedelta(seconds=60), + ), + # A single-active-row task: must be left untouched. + ScheduledTaskRunRow( + id="run-solo", + task_id="task-solo", + thread_id="thread-solo", + scheduled_for=base, + trigger="scheduled", + status="running", + created_at=base, + ), + # A terminal row: must stay terminal. + ScheduledTaskRunRow( + id="run-done", + task_id="task-done", + thread_id="thread-done", + scheduled_for=base, + trigger="scheduled", + status="success", + created_at=base, + ), + ] + ) + session.commit() + finally: + sync_engine.dispose() + + +def _fetch_runs(db_path: Path) -> dict[str, tuple[str, str | None, str | None]]: + """Map id -> (status, error, finished_at) for assertions.""" + with sqlite3.connect(db_path) as raw: + rows = raw.execute("SELECT id, status, error, finished_at FROM scheduled_task_runs").fetchall() + return {run_id: (status, error, finished_at) for run_id, status, error, finished_at in rows} + + +def _index_exists(db_path: Path, index_name: str) -> bool: + with sqlite3.connect(db_path) as raw: + row = raw.execute( + "SELECT 1 FROM sqlite_master WHERE type='index' AND name=?", + (index_name,), + ).fetchone() + return row is not None + + +async def test_migration_supersedes_duplicate_active_runs_before_unique_index(tmp_path: Path) -> None: + db_path = tmp_path / "dirty.db" + _seed_pre_0007_with_duplicates(db_path) + + url = f"sqlite+aiosqlite:///{db_path.as_posix()}" + await init_engine(backend="sqlite", url=url, sqlite_dir=str(tmp_path)) + + try: + runs = _fetch_runs(db_path) + + # Newest active row on the duplicated task survives unchanged. + assert runs["run-newest"][0] == "queued" + assert runs["run-newest"][1] is None + + # Older duplicate active rows are superseded with an explanatory error + # and a finished_at timestamp (mark_stale_active_runs orphan semantics). + for run_id in ("run-old-a", "run-old-b"): + status, error, finished_at = runs[run_id] + assert status == "interrupted", (run_id, status) + assert "uq_scheduled_task_run_active" in (error or ""), (run_id, error) + assert finished_at is not None, run_id + + # Untouched tasks: single active row stays active, terminal stays terminal. + assert runs["run-solo"][0] == "running" + assert runs["run-done"][0] == "success" + + # The partial unique index was successfully created — the upgrade did + # not abort with ``UNIQUE constraint failed``. + assert _index_exists(db_path, "uq_scheduled_task_run_active") + + with sqlite3.connect(db_path) as raw: + version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() + assert version_row[0] == "0007_scheduled_run_active_index" + + # Sanity: the invariant the index enforces now holds — at most one + # active row per task_id. + with sqlite3.connect(db_path) as raw: + dupes = raw.execute("SELECT task_id, COUNT(*) FROM scheduled_task_runs WHERE status IN ('queued', 'running') GROUP BY task_id HAVING COUNT(*) > 1").fetchall() + assert dupes == [] + finally: + await close_engine() diff --git a/backend/tests/test_persistence_bootstrap.py b/backend/tests/test_persistence_bootstrap.py index ca9d84245..537c6cf78 100644 --- a/backend/tests/test_persistence_bootstrap.py +++ b/backend/tests/test_persistence_bootstrap.py @@ -48,7 +48,7 @@ from deerflow.persistence.migrations._helpers import _normalize_default asyncio_test = pytest.mark.asyncio -HEAD = "0006_agents" +HEAD = "0007_scheduled_run_active_index" BASELINE = "0001_baseline" diff --git a/backend/tests/test_persistence_bootstrap_concurrency.py b/backend/tests/test_persistence_bootstrap_concurrency.py index 1848da0f5..aa85acb0d 100644 --- a/backend/tests/test_persistence_bootstrap_concurrency.py +++ b/backend/tests/test_persistence_bootstrap_concurrency.py @@ -28,7 +28,7 @@ from deerflow.persistence.bootstrap import bootstrap_schema pytestmark = pytest.mark.asyncio -HEAD = "0006_agents" +HEAD = "0007_scheduled_run_active_index" def _url(tmp_path: Path) -> str: diff --git a/backend/tests/test_persistence_bootstrap_regression.py b/backend/tests/test_persistence_bootstrap_regression.py index 745fb4773..edec80822 100644 --- a/backend/tests/test_persistence_bootstrap_regression.py +++ b/backend/tests/test_persistence_bootstrap_regression.py @@ -76,7 +76,7 @@ async def test_legacy_database_recovers_token_usage_column(tmp_path: Path) -> No cols = {row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()} assert "token_usage_by_model" in cols version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0006_agents" + assert version_row[0] == "0007_scheduled_run_active_index" # And the read path that originally 500'd must now succeed. sf = get_session_factory() @@ -116,6 +116,6 @@ async def test_legacy_database_with_manual_alter_still_bootstraps(tmp_path: Path # No duplicate column -- list, not set, to catch dupes. assert cols.count("token_usage_by_model") == 1 version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() - assert version_row[0] == "0006_agents" + assert version_row[0] == "0007_scheduled_run_active_index" finally: await close_engine() diff --git a/backend/tests/test_scheduled_task_dispatch_race.py b/backend/tests/test_scheduled_task_dispatch_race.py new file mode 100644 index 000000000..6dcc01dc0 --- /dev/null +++ b/backend/tests/test_scheduled_task_dispatch_race.py @@ -0,0 +1,219 @@ +"""Concurrency regression tests for the scheduled-task dispatch TOCTOU. + +``ScheduledTaskService.dispatch_task`` guards the "at most one active run per +task when overlap_policy=skip" invariant with a non-atomic +``has_active_runs`` check followed by a separate ``create(status="queued")`` +insert. Two concurrent dispatches (double-click, client retry, or a manual +trigger racing the poller) can both pass the check and both launch. The fix +makes the database the atomic arbiter via the partial unique index +``uq_scheduled_task_run_active`` (``task_id WHERE status IN +('queued','running')``); the losing insert is translated to the typed +``ActiveScheduledRunConflict`` and collapsed to the same outcome as the +fast-path check. + +These tests drive the REAL ``ScheduledTaskRunRepository`` + ``ScheduledTaskService`` +against a real file-backed ``sqlite+aiosqlite`` DB (so the index is actually +enforced), with a fake ``launch_run`` that only records launches. +""" + +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime + +import pytest + +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 ActiveScheduledRunConflict, ScheduledTaskRunRepository +from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository + +pytestmark = pytest.mark.asyncio + +_ACTIVE_STATUSES = {"queued", "running"} + + +class _BarrierRunRepo(ScheduledTaskRunRepository): + """Real repository that only releases both dispatchers past + ``has_active_runs`` once both have read it, so their ``create()`` calls + genuinely race for the task's single active slot — a deterministic + reproduction of the check-then-insert TOCTOU.""" + + def __init__(self, session_factory, barrier: asyncio.Barrier | None) -> None: + super().__init__(session_factory) + self._barrier = barrier + + async def has_active_runs(self, task_id: str) -> bool: + result = await super().has_active_runs(task_id) + if self._barrier is not None: + await self._barrier.wait() + return result + + +def _make_service(task_repo, run_repo, launched: list) -> ScheduledTaskService: + async def fake_launch(**kwargs): + # Yield so a truly-concurrent sibling can interleave, then record. + await asyncio.sleep(0) + launched.append(kwargs) + return {"run_id": f"run-{len(launched)}", "thread_id": kwargs["thread_id"]} + + return ScheduledTaskService( + task_repo=task_repo, + task_run_repo=run_repo, + launch_run=fake_launch, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_runs=10, + ) + + +async def _seed_task(task_repo: ScheduledTaskRepository, task_id: str) -> dict: + # fresh_thread_per_run: every dispatch gets a NEW thread_id, so #4003's + # per-thread uq_runs_thread_active can never fire for two dispatches of the + # same task — this is precisely the gap the per-task index closes. + await task_repo.create( + task_id=task_id, + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title=task_id, + prompt="do the thing", + schedule_type="cron", + schedule_spec={"cron": "*/5 * * * *"}, + timezone="UTC", + next_run_at=None, + ) + task = await task_repo.get(task_id, user_id="user-1") + assert task is not None + assert task["overlap_policy"] == "skip" + return task + + +async def _active_run_count(run_repo: ScheduledTaskRunRepository, task_id: str) -> int: + rows = await run_repo.list_by_task(task_id, limit=100) + return sum(1 for row in rows if row["status"] in _ACTIVE_STATUSES) + + +async def test_two_concurrent_manual_dispatches_launch_exactly_once(tmp_path): + 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 = _BarrierRunRepo(sf, asyncio.Barrier(2)) + launched: list = [] + service = _make_service(task_repo, run_repo, launched) + task = await _seed_task(task_repo, "task-race-manual") + now = datetime.now(UTC) + + results = await asyncio.gather( + service.dispatch_task(dict(task), now=now, trigger="manual"), + service.dispatch_task(dict(task), now=now, trigger="manual"), + ) + + outcomes = sorted(result["outcome"] for result in results) + # Exactly one wins the active slot; the loser is a 409-style conflict. + assert outcomes == ["conflict", "launched"], outcomes + assert len(launched) == 1, launched + assert await _active_run_count(run_repo, "task-race-manual") == 1 + # The manual loser records no run-history row (nothing was scheduled). + conflict = next(r for r in results if r["outcome"] == "conflict") + assert conflict["task_run_id"] is None + finally: + await close_engine() + + +async def test_scheduled_and_manual_dispatch_launch_exactly_once(tmp_path): + 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 = _BarrierRunRepo(sf, asyncio.Barrier(2)) + launched: list = [] + service = _make_service(task_repo, run_repo, launched) + task = await _seed_task(task_repo, "task-race-mixed") + now = datetime.now(UTC) + + results = await asyncio.gather( + service.dispatch_task(dict(task), now=now, trigger="scheduled"), + service.dispatch_task(dict(task), now=now, trigger="manual"), + ) + + outcomes = sorted(result["outcome"] for result in results) + # Whichever won launched; the loser is conflict (manual) or skipped + # (scheduled). Which one wins is timing-dependent, but exactly one runs. + assert outcomes.count("launched") == 1, outcomes + assert set(outcomes) <= {"launched", "conflict", "skipped"}, outcomes + assert len(launched) == 1, launched + assert await _active_run_count(run_repo, "task-race-mixed") == 1 + finally: + await close_engine() + + +async def test_natural_timing_concurrent_dispatch_launches_exactly_once(tmp_path): + # No barrier: exercise the fix under the same natural interleaving that + # reproduced the bug (5/5 both-launch on main). The fix must hold whether + # the second dispatch is caught by the has_active_runs fast path or by the + # index-violation path. + 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) + for i in range(5): + launched: list = [] + service = _make_service(task_repo, run_repo, launched) + task_id = f"task-natural-{i}" + task = await _seed_task(task_repo, task_id) + now = datetime.now(UTC) + + results = await asyncio.gather( + service.dispatch_task(dict(task), now=now, trigger="manual"), + service.dispatch_task(dict(task), now=now, trigger="manual"), + ) + + outcomes = sorted(result["outcome"] for result in results) + assert outcomes.count("launched") == 1, (i, outcomes) + assert len(launched) == 1, (i, launched) + assert await _active_run_count(run_repo, task_id) == 1, i + finally: + await close_engine() + + +async def test_partial_unique_index_enforces_one_active_run_per_task(tmp_path): + # Focused repository-level test of the index semantics + the typed conflict. + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + run_repo = ScheduledTaskRunRepository(sf) + now = datetime(2026, 7, 2, 1, 0, tzinfo=UTC) + + await run_repo.create(run_record_id="r1", task_id="t1", thread_id="th1", scheduled_for=now, trigger="scheduled", status="queued") + + # queued -> running is a same-row UPDATE: keeps the one active slot, no + # violation (this is the normal launch transition). + await run_repo.update_status("r1", status="running", run_id="run-1", started_at=now) + assert await run_repo.has_active_runs("t1") is True + + # A second active insert for the same task is a domain conflict. + with pytest.raises(ActiveScheduledRunConflict): + await run_repo.create(run_record_id="r2", task_id="t1", thread_id="th2", scheduled_for=now, trigger="manual", status="queued") + + # Terminal-status rows for the same task are outside the index predicate. + await run_repo.create(run_record_id="r3", task_id="t1", thread_id="th3", scheduled_for=now, trigger="scheduled", status="skipped") + + # A different task's active row is independent. + await run_repo.create(run_record_id="r4", task_id="t2", thread_id="th4", scheduled_for=now, trigger="scheduled", status="queued") + + # Finishing the active run frees the slot; a fresh active row is allowed. + await run_repo.update_status("r1", status="success", run_id="run-1", finished_at=now) + assert await run_repo.has_active_runs("t1") is False + await run_repo.create(run_record_id="r5", task_id="t1", thread_id="th5", scheduled_for=now, trigger="scheduled", status="queued") + assert await run_repo.has_active_runs("t1") is True + finally: + await close_engine() diff --git a/backend/tests/test_scheduled_task_repository.py b/backend/tests/test_scheduled_task_repository.py index 77b5a0710..a9101f440 100644 --- a/backend/tests/test_scheduled_task_repository.py +++ b/backend/tests/test_scheduled_task_repository.py @@ -67,6 +67,10 @@ async def test_mark_stale_active_runs_fails_orphaned_runs(tmp_path): assert sf is not None repo = ScheduledTaskRunRepository(sf) + # The queued and running rows live on different tasks: mark_stale_active_runs + # is a global sweep (no task filter), and the uq_scheduled_task_run_active + # partial unique index forbids two active rows on one task_id, so the pair + # that proves both active statuses get swept must be spread across tasks. await repo.create( run_record_id="task-run-queued", task_id="task-1", @@ -77,12 +81,14 @@ async def test_mark_stale_active_runs_fails_orphaned_runs(tmp_path): ) await repo.create( run_record_id="task-run-running", - task_id="task-1", - thread_id="thread-1", + task_id="task-2", + thread_id="thread-2", scheduled_for=datetime(2026, 7, 2, 1, 0, tzinfo=UTC), trigger="scheduled", status="running", ) + # A terminal row on task-1 (outside the index predicate) coexists with the + # active queued row and must be left untouched by the sweep. await repo.create( run_record_id="task-run-success", task_id="task-1", @@ -95,8 +101,8 @@ async def test_mark_stale_active_runs_fails_orphaned_runs(tmp_path): swept = await repo.mark_stale_active_runs(error="interrupted: gateway restarted") assert swept == 2 - history = await repo.list_by_task("task-1") - by_id = {entry["id"]: entry for entry in history} + by_id = {entry["id"]: entry for entry in await repo.list_by_task("task-1")} + by_id.update({entry["id"]: entry for entry in await repo.list_by_task("task-2")}) assert by_id["task-run-queued"]["status"] == "interrupted" assert by_id["task-run-running"]["status"] == "interrupted" assert by_id["task-run-success"]["status"] == "success" diff --git a/backend/tests/test_scheduled_task_service.py b/backend/tests/test_scheduled_task_service.py index e837f3bb6..07ee0b8a9 100644 --- a/backend/tests/test_scheduled_task_service.py +++ b/backend/tests/test_scheduled_task_service.py @@ -447,7 +447,11 @@ async def test_skip_policy_applies_to_fresh_thread_runs(): assert result["outcome"] == "skipped" assert launched == [] - assert run_repo.created["status"] == "queued" + # The skip tombstone is created directly as terminal "skipped" (not the + # transient "queued" the launch path uses): a queued row is active and would + # itself trip the uq_scheduled_task_run_active partial unique index against + # the pre-existing run still holding the task's single active slot. + assert run_repo.created["status"] == "skipped" assert run_repo.updated[-1][1]["status"] == "skipped" assert task_repo.updated[1]["status"] == "enabled" assert task_repo.updated[1]["increment_run_count"] is False