diff --git a/backend/app/gateway/deps.py b/backend/app/gateway/deps.py index 93197f2ff..a41b7d8fa 100644 --- a/backend/app/gateway/deps.py +++ b/backend/app/gateway/deps.py @@ -47,9 +47,17 @@ _RUN_DRAIN_TIMEOUT_SECONDS = 5.0 def _enforce_postgres_for_multi_worker(config: AppConfig) -> None: - """Refuse to start when GATEWAY_WORKERS > 1 and the DB backend is not Postgres. + """Refuse to start when GATEWAY_WORKERS > 1 and safety preconditions are not met. + + Two checks (both must pass for multi-worker): + + 1. The DB backend must be Postgres — SQLite write-locks cannot support + concurrent multi-process access. + 2. ``run_ownership.heartbeat_enabled`` must be True — without heartbeat, + every run has a NULL lease, so reconciliation treats all inflight + runs as orphans and Worker B would kill Worker A's live runs on + every rolling update or scale-up. - SQLite write-locks cannot support concurrent multi-process access. This gate runs once at startup before any persistence engine is initialised so the error message is clear and the process exits immediately. @@ -66,6 +74,16 @@ def _enforce_postgres_for_multi_worker(config: AppConfig) -> None: if backend != "postgres": raise SystemExit(f"GATEWAY_WORKERS={workers} requires database.backend='postgres', but database.backend is '{backend}'. SQLite cannot support concurrent multi-process access. Set GATEWAY_WORKERS=1 or switch to Postgres.") + run_ownership = getattr(config, "run_ownership", None) + if run_ownership is None or not run_ownership.heartbeat_enabled: + raise SystemExit( + f"GATEWAY_WORKERS={workers} requires run_ownership.heartbeat_enabled=true. " + "Without heartbeat, every run has a NULL lease, so reconciliation " + "treats all inflight runs as orphans — Worker B would kill Worker A's " + "live runs on every rolling update or scale-up. " + "Set run_ownership.heartbeat_enabled=true in config.yaml." + ) + async def _drain_inflight_runs(run_manager: RunManager) -> None: """Drain in-flight runs before the checkpointer is torn down (issue #3373). @@ -287,20 +305,29 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen app.state.run_event_store = make_run_event_store(run_events_config) # RunManager with store backing for persistence - app.state.run_manager = RunManager(store=app.state.run_store) - if getattr(config.database, "backend", None) == "sqlite": - from deerflow.utils.time import now_iso + run_ownership_config = getattr(config, "run_ownership", None) + app.state.run_manager = RunManager( + store=app.state.run_store, + run_ownership_config=run_ownership_config, + ) + # Startup recovery: mark inflight runs whose lease has expired as error. + # In single-worker mode (SQLite / backend=memory), no run has a lease, so + # all inflight rows are reclaimed (unchanged behaviour). In multi-worker + # mode (Postgres), only runs with an expired lease are reclaimed; runs + # owned by another live worker are skipped. + from deerflow.utils.time import now_iso - # Startup-only recovery: clean shutdowns return no active rows and - # the thread-status update below becomes a no-op. - recovered_runs = await app.state.run_manager.reconcile_orphaned_inflight_runs( - error="Gateway restarted before this run reached a durable final state.", - before=now_iso(), - ) - sb_config = getattr(config, "stream_bridge", None) - cleanup_delay = getattr(sb_config, "recovered_stream_cleanup_delay_seconds", 60.0) if sb_config else 60.0 - await _publish_recovered_run_stream_end(app.state.stream_bridge, recovered_runs, cleanup_delay=cleanup_delay) - await _mark_latest_recovered_threads_error(app.state.run_manager, app.state.thread_store, recovered_runs) + recovered_runs = await app.state.run_manager.reconcile_orphaned_inflight_runs( + error="Gateway restarted before this run reached a durable final state.", + before=now_iso(), + ) + sb_config = getattr(config, "stream_bridge", None) + cleanup_delay = getattr(sb_config, "recovered_stream_cleanup_delay_seconds", 60.0) if sb_config else 60.0 + await _publish_recovered_run_stream_end(app.state.stream_bridge, recovered_runs, cleanup_delay=cleanup_delay) + await _mark_latest_recovered_threads_error(app.state.run_manager, app.state.thread_store, recovered_runs) + + # Start the lease heartbeat if enabled (multi-worker deployments). + await app.state.run_manager.start_heartbeat() try: yield diff --git a/backend/packages/harness/deerflow/config/app_config.py b/backend/packages/harness/deerflow/config/app_config.py index 1b53fd30e..b46b0e958 100644 --- a/backend/packages/harness/deerflow/config/app_config.py +++ b/backend/packages/harness/deerflow/config/app_config.py @@ -25,6 +25,7 @@ from deerflow.config.model_config import ModelConfig from deerflow.config.read_before_write_config import ReadBeforeWriteConfig from deerflow.config.reload_boundary import format_field_description from deerflow.config.run_events_config import RunEventsConfig +from deerflow.config.run_ownership_config import RunOwnershipConfig from deerflow.config.runtime_paths import existing_project_file from deerflow.config.safety_finish_reason_config import SafetyFinishReasonConfig from deerflow.config.sandbox_config import SandboxConfig @@ -219,6 +220,13 @@ class AppConfig(BaseModel): field_doc="Stream bridge connecting agent workers to SSE endpoints.", ), ) + run_ownership: RunOwnershipConfig = Field( + default_factory=RunOwnershipConfig, + description=format_field_description( + "run_ownership", + field_doc="Run ownership and lease configuration for multi-worker deployments.", + ), + ) # Name -> config lookup tables, (re)built after validation by # ``_build_name_indexes``. They make ``get_model_config`` / ``get_tool_config`` diff --git a/backend/packages/harness/deerflow/config/reload_boundary.py b/backend/packages/harness/deerflow/config/reload_boundary.py index 3080a1c2a..96131fb00 100644 --- a/backend/packages/harness/deerflow/config/reload_boundary.py +++ b/backend/packages/harness/deerflow/config/reload_boundary.py @@ -68,6 +68,10 @@ STARTUP_ONLY_FIELDS: dict[str, str] = { "ScheduledTaskService is constructed and started once during Gateway lifespan startup; enabled, poll_interval_seconds, lease_seconds, " "and max_concurrent_runs are captured into the service instance and the background poller task is not rebuilt on config.yaml edits." ), + "run_ownership": ( + "RunOwnershipConfig is captured once into RunManager at langgraph_runtime() startup; the lease heartbeat background task is created and " + "started there, and heartbeat_enabled / lease_seconds / grace_seconds are not re-read on config.yaml edits." + ), } diff --git a/backend/packages/harness/deerflow/config/run_ownership_config.py b/backend/packages/harness/deerflow/config/run_ownership_config.py new file mode 100644 index 000000000..a0a5e0389 --- /dev/null +++ b/backend/packages/harness/deerflow/config/run_ownership_config.py @@ -0,0 +1,47 @@ +"""Run ownership configuration for multi-worker deployments.""" + +from __future__ import annotations + +from pydantic import BaseModel, Field + + +class RunOwnershipConfig(BaseModel): + """Per-run ownership and lease configuration. + + When ``heartbeat_enabled`` is True, each worker periodically renews + the lease on its active runs. This is required for multi-worker + deployments to detect orphaned runs from crashed workers. + + Clock-sync assumption + --------------------- + Reconciliation compares another worker's UTC ``lease_expires_at`` against + this worker's ``datetime.now(UTC)``. The only skew budget between two + workers' clocks is ``grace_seconds`` (plus whatever heartbeat slop is + left in the current cycle — at most ``lease_seconds / 3``). Worst case, + if the owning worker's heartbeat is just about to fire, a peer whose + clock is more than ``grace_seconds`` ahead can mis-reclaim a still-live + run as an orphan. + + Operators should ensure worker clocks are synchronised (NTP / chrony / + systemd-timesyncd in K8s nodes) within a few seconds. If the + environment cannot guarantee that, raise ``grace_seconds``; the cost is + longer recovery latency for genuinely dead workers + (``lease_seconds + grace_seconds`` from last heartbeat to reclaim). + """ + + lease_seconds: int = Field( + default=30, + ge=5, + description="Seconds before a run lease expires if not renewed. Heartbeat renews every lease_seconds / 3.", + ) + grace_seconds: int = Field( + default=10, + ge=0, + description=( + "Extra seconds past lease expiry before an orphaned run is reclaimed. Also the clock-skew budget between workers — raise it if worker clocks are not tightly synced; cost is slower recovery of genuinely dead-worker runs." + ), + ) + heartbeat_enabled: bool = Field( + default=False, + description="When True, the worker periodically renews leases on its active runs. Enable for multi-worker deployments (GATEWAY_WORKERS > 1).", + ) diff --git a/backend/packages/harness/deerflow/persistence/migrations/versions/0004_run_ownership.py b/backend/packages/harness/deerflow/persistence/migrations/versions/0004_run_ownership.py new file mode 100644 index 000000000..146291d49 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/versions/0004_run_ownership.py @@ -0,0 +1,134 @@ +"""run ownership. + +Revision ID: 0004_run_ownership +Revises: 0003_scheduled_tasks +Create Date: 2026-07-07 +""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +logger = logging.getLogger(__name__) + +revision: str = "0004_run_ownership" +down_revision: str | Sequence[str] | None = "0003_scheduled_tasks" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def _dedupe_active_runs_per_thread() -> None: + """Cancel superseded active rows so the partial unique index can be built. + + ``uq_runs_thread_active`` enforces at most one pending/running row per + ``thread_id``. A DB that already has two+ active rows for the same thread + (reachable in the field: Postgres deployments had reconciliation skipped + by the old sqlite-only gate, and anyone who ran ``GATEWAY_WORKERS>1`` + before this PR can have duplicates) would fail ``CREATE UNIQUE INDEX`` + and abort the alembic upgrade, blocking gateway startup. + + Keep the newest active row per ``thread_id`` (by ``created_at`` DESC, + ``run_id`` DESC as a deterministic tiebreaker) and mark the rest as + ``error``. Cancelled rows get an explanatory ``error`` string so + operators can see why the run was killed. + """ + bind = op.get_bind() + cancel_message = "cancelled during migration 0004_run_ownership: superseded by a newer active run for the same thread (partial unique index uq_runs_thread_active)" + find_dupe_rows = sa.text( + """ + SELECT run_id, thread_id + FROM runs AS r1 + WHERE r1.status IN ('pending', 'running') + AND EXISTS ( + SELECT 1 FROM runs AS r2 + WHERE r2.thread_id = r1.thread_id + AND r2.status IN ('pending', 'running') + AND r2.run_id <> r1.run_id + AND ( + r2.created_at > r1.created_at + OR (r2.created_at = r1.created_at AND r2.run_id > r1.run_id) + ) + ) + """ + ) + rows = list(bind.execute(find_dupe_rows).fetchall()) + if not rows: + return + for run_id, thread_id in rows: + logger.warning( + "migration 0004_run_ownership: cancelling duplicate active run %s on thread %s", + run_id, + thread_id, + ) + bind.execute( + sa.text( + """ + UPDATE runs + SET status = 'error', + error = :error_message + WHERE status IN ('pending', 'running') + AND EXISTS ( + SELECT 1 FROM runs AS r2 + WHERE r2.thread_id = runs.thread_id + AND r2.status IN ('pending', 'running') + AND r2.run_id <> runs.run_id + AND ( + r2.created_at > runs.created_at + OR (r2.created_at = runs.created_at AND r2.run_id > runs.run_id) + ) + ) + """ + ), + {"error_message": cancel_message}, + ) + + +def upgrade() -> None: + from deerflow.persistence.migrations._helpers import safe_add_column + + safe_add_column("runs", sa.Column("owner_worker_id", sa.String(length=128), nullable=True)) + safe_add_column("runs", sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True)) + + # Idempotent index creation: the legacy 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("runs")} + if "ix_runs_lease" not in existing: + with op.batch_alter_table("runs", schema=None) as batch_op: + batch_op.create_index("ix_runs_lease", ["lease_expires_at"], unique=False) + if "uq_runs_thread_active" not in existing: + # Cancel 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_runs_per_thread() + with op.batch_alter_table("runs", schema=None) as batch_op: + batch_op.create_index( + "uq_runs_thread_active", + ["thread_id"], + unique=True, + sqlite_where=sa.text("status IN ('pending', 'running')"), + postgresql_where=sa.text("status IN ('pending', 'running')"), + ) + + +def downgrade() -> None: + bind = op.get_bind() + insp = sa.inspect(bind) + existing = {ix["name"] for ix in insp.get_indexes("runs")} + if "uq_runs_thread_active" in existing: + with op.batch_alter_table("runs", schema=None) as batch_op: + batch_op.drop_index("uq_runs_thread_active") + if "ix_runs_lease" in existing: + with op.batch_alter_table("runs", schema=None) as batch_op: + batch_op.drop_index("ix_runs_lease") + + from deerflow.persistence.migrations._helpers import safe_drop_column + + safe_drop_column("runs", "lease_expires_at") + safe_drop_column("runs", "owner_worker_id") diff --git a/backend/packages/harness/deerflow/persistence/run/model.py b/backend/packages/harness/deerflow/persistence/run/model.py index 1d5f16f48..19b73e032 100644 --- a/backend/packages/harness/deerflow/persistence/run/model.py +++ b/backend/packages/harness/deerflow/persistence/run/model.py @@ -44,7 +44,25 @@ class RunRow(Base): # Follow-up association follow_up_to_run_id: Mapped[str | None] = mapped_column(String(64)) + # Multi-worker run ownership + owner_worker_id: Mapped[str | None] = mapped_column(String(128), nullable=True) + lease_expires_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)) updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), onupdate=lambda: datetime.now(UTC)) - __table_args__ = (Index("ix_runs_thread_status", "thread_id", "status"),) + __table_args__ = ( + Index("ix_runs_thread_status", "thread_id", "status"), + Index("ix_runs_lease", "lease_expires_at"), + # Cross-process atomicity guarantee: at most one pending/running run per + # thread. 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_runs_thread_active", + "thread_id", + unique=True, + sqlite_where=text("status IN ('pending', 'running')"), + postgresql_where=text("status IN ('pending', 'running')"), + ), + ) diff --git a/backend/packages/harness/deerflow/persistence/run/sql.py b/backend/packages/harness/deerflow/persistence/run/sql.py index 91968c04d..0f1bcf66e 100644 --- a/backend/packages/harness/deerflow/persistence/run/sql.py +++ b/backend/packages/harness/deerflow/persistence/run/sql.py @@ -8,10 +8,10 @@ minutes -- we don't hold connections across long execution. from __future__ import annotations import json -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from typing import Any -from sqlalchemy import select, update +from sqlalchemy import or_, select, update from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from deerflow.persistence.run.model import RunRow @@ -72,7 +72,7 @@ class RunRepository(RunStore): # Convert datetime to ISO string for consistency with MemoryRunStore. # SQLite drops tzinfo on read despite ``DateTime(timezone=True)`` — # ``coerce_iso`` normalizes naive datetimes as UTC. - for key in ("created_at", "updated_at"): + for key in ("created_at", "updated_at", "lease_expires_at"): val = d.get(key) if isinstance(val, datetime): d[key] = coerce_iso(val) @@ -93,6 +93,8 @@ class RunRepository(RunStore): error=None, created_at=None, follow_up_to_run_id=None, + owner_worker_id: str | None = None, + lease_expires_at: str | None = None, ): """Insert or update a run row. @@ -103,6 +105,7 @@ class RunRepository(RunStore): resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.put") now = datetime.now(UTC) created = datetime.fromisoformat(created_at) if created_at else now + lease_dt = datetime.fromisoformat(lease_expires_at) if lease_expires_at else None values = { "thread_id": thread_id, "assistant_id": assistant_id, @@ -114,6 +117,8 @@ class RunRepository(RunStore): "kwargs_json": self._safe_json(kwargs) or {}, "error": error, "follow_up_to_run_id": follow_up_to_run_id, + "owner_worker_id": owner_worker_id, + "lease_expires_at": lease_dt, "updated_at": now, } async with self._sf() as session: @@ -376,3 +381,151 @@ class RunRepository(RunStore): "middleware": middleware, }, } + + # ------------------------------------------------------------------ + # Multi-worker run ownership methods + # ------------------------------------------------------------------ + + async def update_lease( + self, + run_id: str, + *, + owner_worker_id: str, + lease_expires_at: str, + ) -> bool: + lease_dt = datetime.fromisoformat(lease_expires_at) + values: dict[str, Any] = { + "owner_worker_id": owner_worker_id, + "lease_expires_at": lease_dt, + "updated_at": datetime.now(UTC), + } + async with self._sf() as session: + result = await session.execute(update(RunRow).where(RunRow.run_id == run_id, RunRow.owner_worker_id == owner_worker_id, RunRow.status.in_(("pending", "running"))).values(**values)) + await session.commit() + return result.rowcount != 0 + + async def list_inflight_with_expired_lease( + self, + *, + before: str | None = None, + grace_seconds: int = 10, + ) -> list[dict[str, Any]]: + if before is None: + before_dt = datetime.now(UTC) + elif isinstance(before, datetime): + before_dt = before + else: + before_dt = datetime.fromisoformat(before) + cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds) + stmt = ( + select(RunRow) + .where( + RunRow.status.in_(("pending", "running")), + RunRow.created_at <= before_dt, + or_( + RunRow.lease_expires_at.is_(None), + RunRow.lease_expires_at < cutoff, + ), + ) + .order_by(RunRow.created_at.asc()) + ) + async with self._sf() as session: + result = await session.execute(stmt) + return [self._row_to_dict(r) for r in result.scalars()] + + async def create_run_atomic( + self, + run_id: str, + *, + thread_id: str, + owner_worker_id: str, + lease_expires_at: str | None, + multitask_strategy: str = "reject", + assistant_id: str | None = None, + user_id: str | None = None, + model_name: str | None = None, + metadata: dict[str, Any] | None = None, + kwargs: dict[str, Any] | None = None, + created_at: str | None = None, + grace_seconds: int = 10, + ) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Atomically create a run with cross-process thread-uniqueness. + + - For ``reject``: INSERT, let the partial unique index enforce + single-active-run. Returns ``(row_dict, [])`` on success, raises + ``IntegrityError`` on conflict. + - For ``interrupt`` / ``rollback``: SELECT FOR UPDATE inflight + rows for the thread, cancel them (unless their lease is still valid), + then INSERT the new row — all in one transaction. Returns + ``(row_dict, claimed_row_dicts)``. + + Returns: + Tuple of ``(new_run_dict, claimed_run_dicts)``. + """ + from deerflow.runtime.runs.manager import ConflictError + + resolved_user_id = resolve_user_id(user_id or AUTO, method_name="RunRepository.create_run_atomic") + now = datetime.now(UTC) + created = datetime.fromisoformat(created_at) if created_at else now + lease_dt = datetime.fromisoformat(lease_expires_at) if lease_expires_at else None + cutoff = now - timedelta(seconds=grace_seconds) + + values = { + "thread_id": thread_id, + "assistant_id": assistant_id, + "user_id": resolved_user_id, + "model_name": self._normalize_model_name(model_name), + "status": "pending", + "multitask_strategy": multitask_strategy, + "metadata_json": self._safe_json(metadata) or {}, + "kwargs_json": self._safe_json(kwargs) or {}, + "owner_worker_id": owner_worker_id, + "lease_expires_at": lease_dt, + "created_at": created, + "updated_at": now, + } + + async with self._sf() as session: + claimed: list[dict[str, Any]] = [] + + if multitask_strategy in ("interrupt", "rollback"): + stmt = ( + select(RunRow) + .where( + RunRow.thread_id == thread_id, + RunRow.status.in_(("pending", "running")), + ) + .with_for_update() + ) + result = await session.execute(stmt) + for row in result.scalars(): + if row.lease_expires_at is not None: + # SQLite drops tzinfo on read despite + # ``DateTime(timezone=True)`` (see ``_row_to_dict``). + # Treat naive values as UTC — same convention as + # ``coerce_iso`` — so the Python-side comparison + # against the aware ``cutoff`` does not raise + # ``TypeError: can't compare offset-naive and + # offset-aware datetimes`` when heartbeat is enabled + # on SQLite. + row_lease = row.lease_expires_at + if row_lease.tzinfo is None: + row_lease = row_lease.replace(tzinfo=UTC) + if row_lease >= cutoff and row.owner_worker_id != owner_worker_id: + # Live run owned by another worker — we cannot + # interrupt it and the partial unique index would + # reject our INSERT anyway. Surface as + # ConflictError so the caller gets a clean signal + # instead of a retry loop on IntegrityError. + raise ConflictError(f"Thread {thread_id} already has an active run owned by another worker") + row.status = "interrupted" + row.error = "Cancelled by newer run" + row.owner_worker_id = owner_worker_id + row.updated_at = now + claimed.append(self._row_to_dict(row)) + + session.add(RunRow(run_id=run_id, **values)) + await session.commit() + + new_row = await session.get(RunRow, run_id) + return self._row_to_dict(new_row), claimed diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 60646be7b..1715cfda2 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -4,17 +4,22 @@ from __future__ import annotations import asyncio import logging +import socket import sqlite3 import uuid from collections.abc import Awaitable, Callable from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta from typing import TYPE_CHECKING, Any +from sqlalchemy.exc import IntegrityError as SAIntegrityError + from deerflow.utils.time import now_iso as _now_iso from .schemas import DisconnectMode, RunStatus if TYPE_CHECKING: + from deerflow.config.run_ownership_config import RunOwnershipConfig from deerflow.runtime.runs.store.base import RunStore logger = logging.getLogger(__name__) @@ -30,6 +35,72 @@ _RETRYABLE_SQLITE_ERROR_CODES = { sqlite3.SQLITE_LOCKED, } +# Driver-native unique-constraint signals. These are stable across driver and +# SQLAlchemy versions — message text is not (SQLite says "UNIQUE constraint +# failed", Postgres says "duplicate key value violates unique constraint"). +_UNIQUE_PGCODE = "23505" +_SQLITE_UNIQUE_ERRORCODE = sqlite3.SQLITE_CONSTRAINT_UNIQUE + + +def _generate_worker_id() -> str: + """Generate a unique worker identifier: ``hostname:hex_uuid``.""" + return f"{socket.gethostname()}:{uuid.uuid4().hex}" + + +def _is_unique_violation(exc: BaseException) -> bool: + """Return True when *exc* (or its cause chain) is a unique-constraint violation. + + SQLAlchemy wraps the driver's IntegrityError; the wrapped driver exception is + reachable via ``exc.orig`` (and ``__cause__`` / ``__context__``). Prefer + driver-native signals — psycopg ``pgcode`` / ``sqlcode`` = "23505" and + sqlite3 ``sqlite_errorcode`` = ``SQLITE_CONSTRAINT_UNIQUE`` — over message + matching, then fall back to message substrings for cases where the driver + exception isn't reachable through the chain. + + Message text drifts across drivers and locales (SQLite raises + ``UNIQUE constraint failed: .``; Postgres raises + ``duplicate key value violates unique constraint``), so the code/attribute + checks are the load-bearing path. + """ + pending: list[BaseException] = [exc] + seen: set[int] = set() + while pending: + current = pending.pop() + if id(current) in seen: + continue + seen.add(id(current)) + + if getattr(current, "pgcode", None) == _UNIQUE_PGCODE: + return True + if getattr(current, "sqlcode", None) == _UNIQUE_PGCODE: + return True + if getattr(current, "sqlstate", None) == _UNIQUE_PGCODE: + return True + if getattr(current, "sqlite_errorcode", None) == _SQLITE_UNIQUE_ERRORCODE: + return True + + # Message fallbacks are belt-and-suspenders for drivers whose + # native code attribute isn't reachable through the chain. Gate on + # an IntegrityError-typed node so an unrelated application + # exception whose ``str()`` happens to contain "duplicate key" / + # "unique" + "violat" (CHECK constraint message, validation error, + # arbitrary subsystem string) cannot be misclassified as a unique + # violation and silently surface as HTTP 409 instead of 500. + if isinstance(current, (SAIntegrityError, sqlite3.IntegrityError)): + message = str(current).lower() + if "unique constraint failed" in message: + return True + if "unique" in message and "violat" in message: + return True + if "duplicate key" in message: + return True + + for attr in ("orig", "__cause__", "__context__"): + inner = getattr(current, attr, None) + if isinstance(inner, BaseException): + pending.append(inner) + return False + def _is_retryable_persistence_error(exc: BaseException) -> bool: """Return True for transient SQLite persistence failures. @@ -105,6 +176,8 @@ class RunRecord: last_ai_message: str | None = None first_human_message: str | None = None finalizing: bool = False + owner_worker_id: str | None = None + lease_expires_at: str | None = None class RunManager: @@ -120,6 +193,8 @@ class RunManager: store: RunStore | None = None, *, persistence_retry_policy: PersistenceRetryPolicy | None = None, + worker_id: str | None = None, + run_ownership_config: RunOwnershipConfig | None = None, ) -> None: self._runs: dict[str, RunRecord] = {} # Secondary index: thread_id -> insertion-ordered run_id set (a dict is @@ -130,6 +205,10 @@ class RunManager: self._lock = asyncio.Lock() self._store = store self._persistence_retry_policy = persistence_retry_policy or PersistenceRetryPolicy() + self._worker_id = worker_id or _generate_worker_id() + self._run_ownership_config = run_ownership_config + self._heartbeat_task: asyncio.Task | None = None + self._heartbeat_stop: asyncio.Event | None = None def _index_run_locked(self, record: RunRecord) -> None: """Register *record* in the thread index. Caller must hold ``self._lock``.""" @@ -173,6 +252,8 @@ class RunManager: "error": error if error is not None else record.error, "created_at": record.created_at, "model_name": record.model_name, + "owner_worker_id": record.owner_worker_id, + "lease_expires_at": record.lease_expires_at, } if record.user_id is not None: payload["user_id"] = record.user_id @@ -298,6 +379,8 @@ class RunManager: message_count=row.get("message_count") or 0, last_ai_message=row.get("last_ai_message"), first_human_message=row.get("first_human_message"), + owner_worker_id=row.get("owner_worker_id"), + lease_expires_at=row.get("lease_expires_at"), ) async def update_run_completion(self, run_id: str, **kwargs) -> None: @@ -366,9 +449,18 @@ class RunManager: multitask_strategy: str = "reject", user_id: str | None = None, ) -> RunRecord: - """Create a new pending run and register it.""" + """Create a new pending run and register it. + + Note: this method assumes no active run exists for the thread. It + persists via ``store.put`` (upsert) rather than the atomic + ``create_run_atomic`` primitive, so a concurrent insert for the + same thread will hit the partial unique index and surface as a + raw ``IntegrityError`` instead of a ``ConflictError``. Production + callers should use :meth:`create_or_reject`. + """ run_id = str(uuid.uuid4()) now = _now_iso() + lease_expires_at = self._compute_lease_expires_at() record = RunRecord( run_id=run_id, thread_id=thread_id, @@ -381,6 +473,8 @@ class RunManager: user_id=user_id, created_at=now, updated_at=now, + owner_worker_id=self._worker_id, + lease_expires_at=lease_expires_at, ) async with self._lock: self._runs[run_id] = record @@ -601,6 +695,21 @@ class RunManager: logger.info("Run %s cancelled (action=%s)", run_id, action) return True + def _compute_lease_expires_at(self) -> str | None: + """Compute the lease expiration timestamp for a new run. + + Returns ``None`` when heartbeat is disabled (single-worker mode) so + reconciliation treats crashed runs as orphans (NULL lease) and + reclaims them immediately, preserving pre-ownership behaviour. + Multi-worker deployments enable heartbeat, which opts in to leases. + """ + if self._run_ownership_config is None: + return None + if not self._run_ownership_config.heartbeat_enabled: + return None + lease_seconds = self._run_ownership_config.lease_seconds + return (datetime.now(UTC) + timedelta(seconds=lease_seconds)).isoformat() + async def create_or_reject( self, thread_id: str, @@ -619,63 +728,136 @@ class RunManager: already has a pending/running run. For ``interrupt``/``rollback``, cancels inflight runs before creating. - This method holds the lock across both the check and the insert, - eliminating the TOCTOU race in separate ``has_inflight`` + ``create``. + Lock ordering invariant: the local ``self._lock`` is held across + the local check, the store insert, and the local register, so the + store insert can never succeed while a same-worker ConflictError + is about to fire (which would leak a pending row in the store). + Cross-process contention is resolved at the store level via a + partial unique index on ``(thread_id) WHERE status IN + ('pending','running')``. """ run_id = str(uuid.uuid4()) now = _now_iso() _supported_strategies = ("reject", "interrupt", "rollback") + if multitask_strategy not in _supported_strategies: + raise UnsupportedStrategyError(f"Multitask strategy '{multitask_strategy}' is not yet supported. Supported strategies: {', '.join(_supported_strategies)}") + + lease_expires_at = self._compute_lease_expires_at() + grace_seconds = self._run_ownership_config.grace_seconds if self._run_ownership_config else 10 + interrupted_records: list[RunRecord] = [] + record = RunRecord( + run_id=run_id, + thread_id=thread_id, + assistant_id=assistant_id, + status=RunStatus.pending, + on_disconnect=on_disconnect, + multitask_strategy=multitask_strategy, + metadata=metadata or {}, + kwargs=kwargs or {}, + user_id=user_id, + created_at=now, + updated_at=now, + model_name=model_name, + owner_worker_id=self._worker_id, + lease_expires_at=lease_expires_at, + ) async with self._lock: - if multitask_strategy not in _supported_strategies: - raise UnsupportedStrategyError(f"Multitask strategy '{multitask_strategy}' is not yet supported. Supported strategies: {', '.join(_supported_strategies)}") + # 1) Local inflight check (same-worker guard; cross-worker is the + # store's partial unique index below). + local_inflight = [r for r in self._thread_records_locked(thread_id) if r.status in (RunStatus.pending, RunStatus.running) or r.finalizing] - inflight = [r for r in self._thread_records_locked(thread_id) if r.status in (RunStatus.pending, RunStatus.running) or r.finalizing] - - if multitask_strategy == "reject" and inflight: + if multitask_strategy == "reject" and local_inflight: raise ConflictError(f"Thread {thread_id} already has an active run") - if multitask_strategy in ("interrupt", "rollback") and inflight: + if multitask_strategy in ("interrupt", "rollback") and local_inflight: logger.info( "Preparing to cancel %d inflight run(s) on thread %s (strategy=%s)", - len(inflight), + len(local_inflight), thread_id, multitask_strategy, ) - record = RunRecord( - run_id=run_id, - thread_id=thread_id, - assistant_id=assistant_id, - status=RunStatus.pending, - on_disconnect=on_disconnect, - multitask_strategy=multitask_strategy, - metadata=metadata or {}, - kwargs=kwargs or {}, - user_id=user_id, - created_at=now, - updated_at=now, - model_name=model_name, - ) + # 2) Persist to store while still holding the local lock. The + # store is the source of truth for cross-process atomicity. + if self._store is not None: + if multitask_strategy == "reject": + try: + await self._call_store_with_retry( + "create_run_atomic", + run_id, + lambda: self._store.create_run_atomic( + run_id=run_id, + thread_id=thread_id, + owner_worker_id=self._worker_id, + lease_expires_at=lease_expires_at, + multitask_strategy="reject", + assistant_id=assistant_id, + user_id=user_id, + model_name=model_name, + metadata=metadata, + kwargs=kwargs, + created_at=now, + grace_seconds=grace_seconds, + ), + ) + except ConflictError: + raise + except Exception as exc: + if _is_unique_violation(exc): + raise ConflictError(f"Thread {thread_id} already has an active run") from exc + raise + else: + # Interrupt / rollback: store-side claim + insert in one + # transaction. Retry on IntegrityError in case another + # worker races us between our SELECT FOR UPDATE and INSERT. + max_retries = 3 + for attempt in range(max_retries): + try: + await self._call_store_with_retry( + "create_run_atomic", + run_id, + lambda: self._store.create_run_atomic( + run_id=run_id, + thread_id=thread_id, + owner_worker_id=self._worker_id, + lease_expires_at=lease_expires_at, + multitask_strategy=multitask_strategy, + assistant_id=assistant_id, + user_id=user_id, + model_name=model_name, + metadata=metadata, + kwargs=kwargs, + created_at=now, + grace_seconds=grace_seconds, + ), + ) + break + except Exception as exc: + is_unique = _is_unique_violation(exc) + if is_unique and attempt + 1 < max_retries: + continue + if is_unique: + # Exhausted retries on unique violation — surface + # as ConflictError to match the reject branch's + # contract (409, not 500). Same root cause: another + # worker won the race for this thread. + raise ConflictError(f"Thread {thread_id} already has an active run") from exc + raise + # ``create_run_atomic`` already marked any claimed store + # rows as interrupted in the same transaction; no extra + # store write is needed for them. + + # 3) Only now safe to register locally — store insert succeeded. self._runs[run_id] = record self._index_run_locked(record) - persisted = False - try: - await self._persist_new_run_to_store(record) - persisted = True - except Exception: - logger.warning("Failed to persist run %s; rolled back in-memory record", run_id, exc_info=True) - raise - finally: - # Also covers cancellation, which bypasses ``except Exception``. - if not persisted: - self._runs.pop(run_id, None) - self._unindex_run_locked(run_id, record.thread_id) - if multitask_strategy in ("interrupt", "rollback") and inflight: - for r in inflight: + # 4) Cancel local in-memory inflight (interrupt/rollback). The + # store-side counterparts were already cancelled in step 2. + if multitask_strategy in ("interrupt", "rollback"): + for r in local_inflight: if r.finalizing: continue r.abort_action = multitask_strategy @@ -688,8 +870,11 @@ class RunManager: r.updated_at = now interrupted_records.append(r) + # Outside the lock: persist interrupted status for locally-cancelled + # runs. Store-side claimed rows are already finalised. for interrupted_record in interrupted_records: await self._persist_status(interrupted_record, RunStatus.interrupted) + logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id) return record @@ -699,22 +884,25 @@ class RunManager: error: str, before: str | None = None, ) -> list[RunRecord]: - """Mark persisted active runs as failed when no local task owns them. + """Mark persisted active runs as failed when their lease has expired. - Gateway runs are process-local: the asyncio task and abort event live in - memory, while the run row is durable. After a SQLite-backed gateway - restart, any persisted ``pending`` or ``running`` row created before - startup cannot still have a local worker. This recovery step turns that - ambiguous state into an explicit error instead of letting the UI show an - indefinite active run. + In multi-worker deployments (Postgres), a run owned by Worker A that + still shows ``pending`` / ``running`` after its lease expired means + Worker A crashed or was partitioned. This worker (B) can safely claim + and error it out because the lease was not renewed. + + Rows with a still-valid lease are skipped — they belong to another live + worker. Rows with a NULL lease (pre-ownership data) are reclaimed as + well, matching the original single-worker recovery behaviour. """ if self._store is None: return [] + grace_seconds = self._run_ownership_config.grace_seconds if self._run_ownership_config else 10 try: rows = await self._call_store_with_retry( - "list_inflight", + "list_inflight_with_expired_lease", "*", - lambda: self._store.list_inflight(before=before), + lambda: self._store.list_inflight_with_expired_lease(before=before, grace_seconds=grace_seconds), ) except Exception: logger.warning("Failed to list orphaned inflight runs for reconciliation", exc_info=True) @@ -732,6 +920,7 @@ class RunManager: async with self._lock: live_record = self._runs.get(record.run_id) if live_record is not None and live_record.status in (RunStatus.pending, RunStatus.running): + # Still owned by a local task — skip continue record.status = RunStatus.error @@ -762,9 +951,162 @@ class RunManager: self._unindex_run_locked(run_id, record.thread_id) logger.debug("Run record %s cleaned up", run_id) + # ------------------------------------------------------------------ + # Lease heartbeat + # ------------------------------------------------------------------ + + @property + def worker_id(self) -> str: + """Return this worker's unique identifier.""" + return self._worker_id + + @property + def heartbeat_enabled(self) -> bool: + """Return ``True`` when the heartbeat background task should run.""" + if self._run_ownership_config is None: + return False + return self._run_ownership_config.heartbeat_enabled + + async def start_heartbeat(self) -> None: + """Start the background lease-renewal task. + + No-op when ``heartbeat_enabled`` is ``False`` or the task is already running. + """ + if not self.heartbeat_enabled: + return + if self._heartbeat_task is not None and not self._heartbeat_task.done(): + return + self._heartbeat_stop = asyncio.Event() + task = asyncio.create_task(self._heartbeat_loop()) + task.set_name("deerflow-run-lease-heartbeat") + self._heartbeat_task = task + logger.info("Run lease heartbeat started for worker %s", self._worker_id) + + async def stop_heartbeat(self) -> None: + """Stop the background heartbeat task.""" + if self._heartbeat_stop is not None: + self._heartbeat_stop.set() + if self._heartbeat_task is not None and not self._heartbeat_task.done(): + try: + await asyncio.wait_for(self._heartbeat_task, timeout=5.0) + except TimeoutError: + self._heartbeat_task.cancel() + try: + await self._heartbeat_task + except asyncio.CancelledError: + pass + except asyncio.CancelledError: + pass + self._heartbeat_task = None + self._heartbeat_stop = None + logger.info("Run lease heartbeat stopped for worker %s", self._worker_id) + + async def _heartbeat_loop(self) -> None: + """Periodically renew leases and reclaim orphaned runs from dead peers. + + Lease renewal runs every ``lease_seconds / 3``. Reconciliation + (sweeping for expired leases owned by dead workers) runs every + ``lease_seconds`` (every 3rd cycle) so orphaned runs are recovered + without waiting for a pod restart. + + Both operations are guarded so a transient failure cannot take the + heartbeat task down — a dead heartbeat means no lease is renewed + again, and every active run eventually looks orphaned to peers. + """ + if self._run_ownership_config is None or self._heartbeat_stop is None: + return + lease_seconds = self._run_ownership_config.lease_seconds + interval = max(1, lease_seconds // 3) + stop = self._heartbeat_stop + cycle = 0 + + while not stop.is_set(): + try: + await asyncio.wait_for(stop.wait(), timeout=interval) + break # stop event was set + except TimeoutError: + pass # interval elapsed + + cycle += 1 + try: + await self._renew_leases() + except Exception: + logger.warning("Heartbeat renewal cycle failed", exc_info=True) + + # Reconcile every 3rd cycle (= every lease_seconds). Startup + # reconciliation (in langgraph_runtime) covers the initial + # sweep; this periodic pass catches orphans whose lease + # expires between restarts — e.g. Worker A crashes, its + # replacement starts before the lease expires, and the + # startup pass skips the still-valid lease. + if cycle % 3 == 0: + try: + await self._reconcile_orphans_periodic() + except Exception: + logger.warning("Periodic orphan reconciliation failed", exc_info=True) + + async def _renew_leases(self) -> None: + """Renew the lease on every locally-owned active run.""" + if self._store is None or self._run_ownership_config is None: + return + lease_seconds = self._run_ownership_config.lease_seconds + new_expiry = (datetime.now(UTC) + timedelta(seconds=lease_seconds)).isoformat() + + async with self._lock: + # Renew any pending/running run owned by this worker unless its + # background task has already completed. A pending run whose task + # has not been spawned yet (``task is None``) is still live from + # this worker's perspective — between ``create_run_atomic`` + # inserting the row and the worker layer spawning the agent task + # there is a brief window. If we drop those records here and the + # window stretches past ``lease_seconds`` (e.g. event-loop + # saturation, slow checkpoint hydrate on a fresh worker), peer + # reconciliation will reclaim the run as an orphan and mark it + # ``error`` even though this worker still intends to execute it. + active_runs = [(rid, record) for rid, record in self._runs.items() if record.status in (RunStatus.pending, RunStatus.running) and record.owner_worker_id == self._worker_id and (record.task is None or not record.task.done())] + + for run_id, record in active_runs: + try: + updated = await self._call_store_with_retry( + "update_lease", + run_id, + lambda: self._store.update_lease( + run_id, + owner_worker_id=self._worker_id, + lease_expires_at=new_expiry, + ), + ) + if updated: + # Unsynced write is benign: ``lease_expires_at`` is the + # only field on an existing record this path mutates, so + # there is no concurrent writer to race against + # (``set_status`` / ``_persist_status`` touch other + # fields). Re-acquiring ``self._lock`` here would + # serialise against unrelated run mutations for no gain. + record.lease_expires_at = new_expiry + except Exception: + logger.warning("Failed to renew lease for run %s", run_id, exc_info=True) + + async def _reconcile_orphans_periodic(self) -> None: + """Sweep for expired leases owned by dead peers. + + Called from ``_heartbeat_loop`` every ``lease_seconds``. Startup + reconciliation handles the initial sweep; this periodic pass + catches orphans whose lease expires between restarts. + """ + error_msg = "Run lease expired — owning worker is unreachable." + recovered = await self.reconcile_orphaned_inflight_runs(error=error_msg) + if recovered: + logger.warning( + "Periodic reconciliation recovered %d orphaned run(s) as error", + len(recovered), + ) + async def shutdown(self, *, timeout: float = 5.0) -> None: """Cancel and bounded-await all in-flight runs on process shutdown. + Stops the lease heartbeat first so no renewal races against the drain. + Chat runs execute in fire-and-forget background ``asyncio`` tasks that write checkpoints through a shared checkpointer. On shutdown the checkpointer's resources (e.g. the postgres connection pool owned by the @@ -789,6 +1131,7 @@ class RunManager: ``app.gateway.app._SHUTDOWN_HOOK_TIMEOUT_SECONDS``. Runs still active after ``timeout`` are logged and may still race teardown. """ + await self.stop_heartbeat() loop = asyncio.get_running_loop() deadline = loop.time() + timeout diff --git a/backend/packages/harness/deerflow/runtime/runs/store/base.py b/backend/packages/harness/deerflow/runtime/runs/store/base.py index 682adeaea..1ccc1a8cc 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/base.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/base.py @@ -30,6 +30,8 @@ class RunStore(abc.ABC): kwargs: dict[str, Any] | None = None, error: str | None = None, created_at: str | None = None, + owner_worker_id: str | None = None, + lease_expires_at: str | None = None, ) -> None: pass @@ -142,3 +144,48 @@ class RunStore(abc.ABC): by_caller ({lead_agent, subagent, middleware}). """ pass + + @abc.abstractmethod + async def update_lease( + self, + run_id: str, + *, + owner_worker_id: str, + lease_expires_at: str, + ) -> bool: + """Renew the lease on an active run. Returns ``False`` when no row matched.""" + pass + + @abc.abstractmethod + async def list_inflight_with_expired_lease( + self, + *, + before: str | None = None, + grace_seconds: int = 10, + ) -> list[dict[str, Any]]: + """Return active runs whose lease has expired (or is NULL for pre-ownership rows).""" + pass + + @abc.abstractmethod + async def create_run_atomic( + self, + run_id: str, + *, + thread_id: str, + owner_worker_id: str, + lease_expires_at: str | None, + multitask_strategy: str = "reject", + assistant_id: str | None = None, + user_id: str | None = None, + model_name: str | None = None, + metadata: dict[str, Any] | None = None, + kwargs: dict[str, Any] | None = None, + created_at: str | None = None, + grace_seconds: int = 10, + ) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Atomically create a run row with cross-process thread-uniqueness. + + Returns ``(new_run_dict, claimed_run_dicts)``. + Raises ``IntegrityError`` on conflict for ``reject`` strategy. + """ + pass diff --git a/backend/packages/harness/deerflow/runtime/runs/store/memory.py b/backend/packages/harness/deerflow/runtime/runs/store/memory.py index 902f847ee..a2b98fe86 100644 --- a/backend/packages/harness/deerflow/runtime/runs/store/memory.py +++ b/backend/packages/harness/deerflow/runtime/runs/store/memory.py @@ -5,7 +5,7 @@ Equivalent to the original RunManager._runs dict behavior. from __future__ import annotations -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from typing import Any from deerflow.runtime.runs.store.base import RunStore @@ -46,6 +46,8 @@ class MemoryRunStore(RunStore): kwargs=None, error=None, created_at=None, + owner_worker_id=None, + lease_expires_at=None, ): now = datetime.now(UTC).isoformat() self._runs[run_id] = { @@ -61,6 +63,8 @@ class MemoryRunStore(RunStore): "error": error, "created_at": created_at or now, "updated_at": now, + "owner_worker_id": owner_worker_id, + "lease_expires_at": lease_expires_at, } self._index_run(run_id, thread_id) @@ -166,3 +170,157 @@ class MemoryRunStore(RunStore): "middleware": sum(r.get("middleware_tokens", 0) for r in completed), }, } + + # ------------------------------------------------------------------ + # Multi-worker run ownership methods + # ------------------------------------------------------------------ + + async def update_lease( + self, + run_id: str, + *, + owner_worker_id: str, + lease_expires_at: str, + ) -> bool: + run = self._runs.get(run_id) + if run is None: + return False + if run["status"] not in ("pending", "running"): + return False + if run.get("owner_worker_id") != owner_worker_id: + return False + run["owner_worker_id"] = owner_worker_id + run["lease_expires_at"] = lease_expires_at + run["updated_at"] = datetime.now(UTC).isoformat() + return True + + async def list_inflight_with_expired_lease( + self, + *, + before: str | None = None, + grace_seconds: int = 10, + ) -> list[dict[str, Any]]: + now_dt = datetime.fromisoformat(before) if before else datetime.now(UTC) + cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds) + results = [] + for r in self._runs.values(): + if r["status"] not in ("pending", "running"): + continue + created_at = r.get("created_at", "") + if not created_at: + continue + try: + created_dt = datetime.fromisoformat(created_at) + except (ValueError, TypeError): + continue + if created_dt > now_dt: + continue + lease = r.get("lease_expires_at") + if lease is None: + # Pre-ownership rows: no lease means orphaned + results.append(r) + else: + try: + lease_dt = datetime.fromisoformat(lease) + # Treat naive values as UTC — same convention as + # ``coerce_iso`` in the SQL store, so the comparison + # against the aware ``cutoff`` does not raise + # ``TypeError`` when heartbeat is enabled on SQLite + # (which drops tzinfo on read). + if lease_dt.tzinfo is None: + lease_dt = lease_dt.replace(tzinfo=UTC) + if lease_dt < cutoff: + results.append(r) + except (ValueError, TypeError): + results.append(r) + results.sort(key=lambda r: r["created_at"]) + return results + + async def create_run_atomic( + self, + run_id: str, + *, + thread_id: str, + owner_worker_id: str, + lease_expires_at: str | None, + multitask_strategy: str = "reject", + assistant_id: str | None = None, + user_id: str | None = None, + model_name: str | None = None, + metadata: dict[str, Any] | None = None, + kwargs: dict[str, Any] | None = None, + created_at: str | None = None, + grace_seconds: int = 10, + ) -> tuple[dict[str, Any], list[dict[str, Any]]]: + from deerflow.runtime.runs.manager import ConflictError + + now = datetime.now(UTC).isoformat() + cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds) + + # For reject: check if any active run exists + if multitask_strategy == "reject": + for r in self._runs.values(): + if r["thread_id"] == thread_id and r["status"] in ("pending", "running"): + raise ConflictError(f"Thread {thread_id} already has an active run") + + # For interrupt/rollback: claim inflight runs. + # Two-pass so the memory path mirrors the SQL store's transactional + # semantics — if any candidate is a live run owned by another worker + # we must raise ConflictError WITHOUT having already mutated earlier + # candidates. Mutating inline would leave the store in a half- + # interrupted state on raise, diverging from SQL where a raise rolls + # the whole transaction back. + claimed = [] + if multitask_strategy in ("interrupt", "rollback"): + candidates: list[dict[str, Any]] = [] + for r in self._runs.values(): + if r["thread_id"] != thread_id: + continue + if r["status"] not in ("pending", "running"): + continue + existing_lease = r.get("lease_expires_at") + if existing_lease is not None: + try: + lease_dt = datetime.fromisoformat(existing_lease) + # Treat naive values as UTC — same convention as + # the SQL store and ``coerce_iso``, so the + # comparison against the aware ``cutoff`` does not + # raise ``TypeError``. + if lease_dt.tzinfo is None: + lease_dt = lease_dt.replace(tzinfo=UTC) + if lease_dt >= cutoff and r.get("owner_worker_id") != owner_worker_id: + # Live run owned by another worker — cannot + # interrupt, and the partial unique index would + # reject the INSERT anyway. Surface as ConflictError + # so the caller gets a clean signal. Raise before + # any mutation so the store is left untouched. + raise ConflictError(f"Thread {thread_id} already has an active run owned by another worker") + except (ValueError, TypeError): + pass + candidates.append(r) + for r in candidates: + r["status"] = "interrupted" + r["error"] = "Cancelled by newer run" + r["owner_worker_id"] = owner_worker_id + r["updated_at"] = now + claimed.append(r) + + new_row = { + "run_id": run_id, + "thread_id": thread_id, + "assistant_id": assistant_id, + "user_id": user_id, + "model_name": model_name, + "status": "pending", + "multitask_strategy": multitask_strategy, + "metadata": metadata or {}, + "kwargs": kwargs or {}, + "error": None, + "owner_worker_id": owner_worker_id, + "lease_expires_at": lease_expires_at, + "created_at": created_at or now, + "updated_at": now, + } + self._runs[run_id] = new_row + self._index_run(run_id, thread_id) + return new_row, claimed diff --git a/backend/tests/test_gateway_run_recovery.py b/backend/tests/test_gateway_run_recovery.py index e5f300726..df2eb2308 100644 --- a/backend/tests/test_gateway_run_recovery.py +++ b/backend/tests/test_gateway_run_recovery.py @@ -29,8 +29,9 @@ class _FakeRunManager: recovered_runs = [SimpleNamespace(run_id="run-1", thread_id="thread-1")] latest_by_thread: dict[str, list[SimpleNamespace]] = {} - def __init__(self, *, store): + def __init__(self, *, store, run_ownership_config=None): self.store = store + self.run_ownership_config = run_ownership_config self.reconcile_calls: list[dict] = [] self.list_by_thread_calls: list[dict] = [] self.shutdown_calls: int = 0 @@ -44,6 +45,12 @@ class _FakeRunManager: self.list_by_thread_calls.append({"thread_id": thread_id, "user_id": user_id, "limit": limit}) return self.latest_by_thread.get(thread_id, self.recovered_runs[:limit]) + async def start_heartbeat(self) -> None: + pass + + async def stop_heartbeat(self) -> None: + pass + async def shutdown(self, *, timeout: float = 5.0) -> None: # No in-flight tasks in these startup-recovery tests; langgraph_runtime # drains the manager on teardown, so the double must accept the call. diff --git a/backend/tests/test_migration_0004_run_ownership_dedupe.py b/backend/tests/test_migration_0004_run_ownership_dedupe.py new file mode 100644 index 000000000..db6aa01eb --- /dev/null +++ b/backend/tests/test_migration_0004_run_ownership_dedupe.py @@ -0,0 +1,167 @@ +"""Regression test for migration ``0004_run_ownership`` dedupe pass. + +End-to-end shape: + +1. Hand-build a SQLite DB that mirrors a real pre-0004 deployment that ran + ``GATEWAY_WORKERS>1`` before this PR and accumulated duplicate active rows + per thread (the exact dirty state the multi-worker ownership fix targets). +2. Stamp it at ``0003_scheduled_tasks`` so ``bootstrap_schema`` takes the + versioned branch and runs ``alembic upgrade head``. +3. Insert two+ pending/running rows for the same ``thread_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`` → ``0004.upgrade()``. +5. Verify the migration cancelled the superseded duplicates (set them to + ``error`` with an explanatory message), kept the newest active row, and + successfully built the ``uq_runs_thread_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.run.model import RunRow + +pytestmark = pytest.mark.asyncio + + +def _seed_pre_0004_with_duplicates(db_path: Path) -> None: + """Build a DB at revision 0003 with duplicate active rows per thread. + + Uses a synchronous engine so the seed is independent of the async engine + under test. ``Base.metadata.create_all`` produces the full current schema + (including the partial unique index), so we drop just the unique index to + land in the dirty state the migration's dedupe pass targets: a versioned + DB at 0003 where duplicate active rows per thread can coexist. We then + stamp at 0003 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 — this is the invariant the + # migration rebuilds, and its absence is what permits duplicate + # active rows to exist in the first place. + conn.execute(sa.text("DROP INDEX IF EXISTS uq_runs_thread_active")) + # Stamp at 0003 so bootstrap takes the versioned branch and runs + # ``alembic upgrade head`` (which is what executes 0004.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 ('0003_scheduled_tasks')")) + + base = datetime.now(UTC) + with Session(sync_engine) as session: + session.add_all( + [ + RunRow( + run_id="run-old-a", + thread_id="thread-dup", + status="pending", + created_at=base, + updated_at=base, + ), + RunRow( + run_id="run-old-b", + thread_id="thread-dup", + status="running", + created_at=base + timedelta(seconds=10), + updated_at=base + timedelta(seconds=10), + ), + RunRow( + run_id="run-newest", + thread_id="thread-dup", + status="pending", + created_at=base + timedelta(seconds=60), + updated_at=base + timedelta(seconds=60), + ), + RunRow( + run_id="run-solo", + thread_id="thread-solo", + status="running", + created_at=base, + updated_at=base, + ), + RunRow( + run_id="run-success", + thread_id="thread-done", + status="success", + created_at=base, + updated_at=base, + ), + ] + ) + session.commit() + finally: + sync_engine.dispose() + + +def _fetch_runs(db_path: Path) -> dict[str, tuple[str, str | None]]: + """Map run_id -> (status, error) for assertions.""" + with sqlite3.connect(db_path) as raw: + rows = raw.execute("SELECT run_id, status, error FROM runs").fetchall() + return {run_id: (status, error) for run_id, status, error 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_dedupes_duplicate_active_rows_before_unique_index(tmp_path: Path) -> None: + db_path = tmp_path / "dirty.db" + _seed_pre_0004_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 thread survives unchanged. + assert runs["run-newest"] == ("pending", None) + + # Older duplicate active rows are cancelled with an explanatory error. + assert runs["run-old-a"][0] == "error" + assert "uq_runs_thread_active" in (runs["run-old-a"][1] or "") + assert runs["run-old-b"][0] == "error" + assert "uq_runs_thread_active" in (runs["run-old-b"][1] or "") + + # Untouched threads: single active row stays active, terminal rows stay terminal. + assert runs["run-solo"] == ("running", None) + assert runs["run-success"] == ("success", None) + + # The partial unique index was successfully created — the upgrade did + # not abort with ``UNIQUE constraint failed``. + assert _index_exists(db_path, "uq_runs_thread_active") + assert _index_exists(db_path, "ix_runs_lease") + + with sqlite3.connect(db_path) as raw: + version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone() + assert version_row[0] == "0004_run_ownership" + + # Sanity: the invariant the index enforces is now true — at most one + # active row per thread. + with sqlite3.connect(db_path) as raw: + dupes = raw.execute("SELECT thread_id, COUNT(*) FROM runs WHERE status IN ('pending', 'running') GROUP BY thread_id HAVING COUNT(*) > 1").fetchall() + assert dupes == [] + finally: + await close_engine() diff --git a/backend/tests/test_multi_worker_postgres_gate.py b/backend/tests/test_multi_worker_postgres_gate.py index f87e5e7b9..0a059ba6f 100644 --- a/backend/tests/test_multi_worker_postgres_gate.py +++ b/backend/tests/test_multi_worker_postgres_gate.py @@ -20,10 +20,12 @@ from fastapi import FastAPI from app.gateway.deps import _enforce_postgres_for_multi_worker, langgraph_runtime from deerflow.config.database_config import DatabaseConfig +from deerflow.config.run_ownership_config import RunOwnershipConfig -def _config_with_backend(backend: str) -> SimpleNamespace: - return SimpleNamespace(database=DatabaseConfig(backend=backend)) +def _config_with_backend(backend: str, *, heartbeat_enabled: bool | None = None) -> SimpleNamespace: + run_ownership = RunOwnershipConfig(heartbeat_enabled=heartbeat_enabled) if heartbeat_enabled is not None else None + return SimpleNamespace(database=DatabaseConfig(backend=backend), run_ownership=run_ownership) # --------------------------------------------------------------------------- @@ -45,9 +47,9 @@ def test_gate_noop_for_single_worker(monkeypatch): _enforce_postgres_for_multi_worker(_config_with_backend(backend)) -def test_gate_allows_multi_worker_with_postgres(monkeypatch): +def test_gate_allows_multi_worker_with_postgres_and_heartbeat(monkeypatch): monkeypatch.setenv("GATEWAY_WORKERS", "2") - _enforce_postgres_for_multi_worker(_config_with_backend("postgres")) + _enforce_postgres_for_multi_worker(_config_with_backend("postgres", heartbeat_enabled=True)) def test_gate_rejects_multi_worker_with_sqlite(monkeypatch): @@ -103,6 +105,42 @@ def test_gate_error_message_lists_both_remediations(monkeypatch): assert "Postgres" in msg, "must mention the alternative backend" +# --------------------------------------------------------------------------- +# Heartbeat enforcement: multi-worker requires heartbeat_enabled=true +# --------------------------------------------------------------------------- + + +def test_gate_rejects_multi_worker_without_heartbeat(monkeypatch): + monkeypatch.setenv("GATEWAY_WORKERS", "2") + with pytest.raises(SystemExit) as exc_info: + _enforce_postgres_for_multi_worker(_config_with_backend("postgres", heartbeat_enabled=False)) + msg = str(exc_info.value) + assert "heartbeat_enabled=true" in msg + + +def test_gate_rejects_multi_worker_without_run_ownership_config(monkeypatch): + monkeypatch.setenv("GATEWAY_WORKERS", "2") + with pytest.raises(SystemExit) as exc_info: + _enforce_postgres_for_multi_worker(_config_with_backend("postgres", heartbeat_enabled=None)) + msg = str(exc_info.value) + assert "heartbeat_enabled=true" in msg + + +def test_gate_heartbeat_check_not_triggered_for_single_worker(monkeypatch): + """GATEWAY_WORKERS=1 skips the heartbeat check entirely.""" + monkeypatch.setenv("GATEWAY_WORKERS", "1") + _enforce_postgres_for_multi_worker(_config_with_backend("postgres", heartbeat_enabled=False)) + + +def test_gate_heartbeat_check_not_triggered_for_sqlite(monkeypatch): + """The gate exits on Postgres check before reaching heartbeat check.""" + monkeypatch.setenv("GATEWAY_WORKERS", "2") + with pytest.raises(SystemExit) as exc_info: + _enforce_postgres_for_multi_worker(_config_with_backend("sqlite", heartbeat_enabled=True)) + msg = str(exc_info.value) + assert "postgres" in msg.lower() + + # --------------------------------------------------------------------------- # Integration: the gate is wired into langgraph_runtime before init_engine # --------------------------------------------------------------------------- diff --git a/backend/tests/test_multi_worker_run_ownership.py b/backend/tests/test_multi_worker_run_ownership.py new file mode 100644 index 000000000..1d16f182c --- /dev/null +++ b/backend/tests/test_multi_worker_run_ownership.py @@ -0,0 +1,967 @@ +"""Tests for multi-worker run ownership (work items 2–3). + +Coverage: +- create_or_reject with reject strategy blocks duplicate active runs +- create_or_reject with interrupt strategy claims and cancels old runs +- create_run_atomic refuses to interrupt a run owned by another live worker +- reconcile_orphaned_inflight_runs uses lease-based detection +- Worker reconciliation skips runs with unexpired leases +- Lease heartbeat renews active run leases +- GATEWAY_WORKERS=1 + heartbeat_enabled=false behaviour unchanged +""" + +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime, timedelta +from unittest.mock import AsyncMock + +import pytest + +from deerflow.config.run_ownership_config import RunOwnershipConfig +from deerflow.runtime import RunManager, RunStatus +from deerflow.runtime.runs.manager import ConflictError, _generate_worker_id +from deerflow.runtime.runs.store.memory import MemoryRunStore + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _lease_config(**kwargs) -> RunOwnershipConfig: + return RunOwnershipConfig( + lease_seconds=kwargs.get("lease_seconds", 30), + grace_seconds=kwargs.get("grace_seconds", 10), + heartbeat_enabled=kwargs.get("heartbeat_enabled", False), + ) + + +def _make_manager(store=None, **kwargs) -> RunManager: + return RunManager( + store=store or MemoryRunStore(), + run_ownership_config=kwargs.pop("run_ownership_config", _lease_config()), + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# create_or_reject — reject strategy +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_reject_blocks_when_active_run_exists(): + """reject strategy must raise ConflictError when thread has an active run.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + await manager.create("thread-1") + await manager.set_status((await manager.list_by_thread("thread-1"))[0].run_id, RunStatus.running) + + with pytest.raises(ConflictError, match="already has an active run"): + await manager.create_or_reject("thread-1", multitask_strategy="reject") + + +@pytest.mark.anyio +async def test_reject_succeeds_when_no_active_run(): + """reject strategy must succeed when the thread has no active run.""" + store = MemoryRunStore() + manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True)) + record = await manager.create_or_reject("thread-1", multitask_strategy="reject") + assert record is not None + assert record.status == RunStatus.pending + assert record.owner_worker_id is not None + assert record.lease_expires_at is not None + + +@pytest.mark.anyio +async def test_reject_blocks_reentrant_same_thread_locally(): + """reject must also block when a local in-memory active run exists.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + await manager.create_or_reject("thread-1", multitask_strategy="reject") + + with pytest.raises(ConflictError, match="already has an active run"): + await manager.create_or_reject("thread-1", multitask_strategy="reject") + + +# --------------------------------------------------------------------------- +# create_or_reject — interrupt strategy +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_interrupt_cancels_old_run_and_creates_new(): + """interrupt must cancel the previous active run and create a new one.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + old = await manager.create_or_reject("thread-1", multitask_strategy="reject") + await manager.set_status(old.run_id, RunStatus.running) + + new = await manager.create_or_reject("thread-1", multitask_strategy="interrupt") + + assert new.run_id != old.run_id + assert new.status == RunStatus.pending + + # Old run must be interrupted locally + assert old.status == RunStatus.interrupted + assert old.abort_event.is_set() + + # Old run must be marked interrupted in-store (persist_status after local cancel) + old_after = await store.get(old.run_id) + assert old_after["status"] == "interrupted" + + +@pytest.mark.anyio +async def test_interrupt_creates_new_when_old_completed(): + """interrupt must succeed when the previous run already reached a terminal status.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + old = await manager.create_or_reject("thread-1") + await manager.set_status(old.run_id, RunStatus.success) + + new = await manager.create_or_reject("thread-1", multitask_strategy="interrupt") + assert new.run_id != old.run_id + assert new.status == RunStatus.pending + + +@pytest.mark.anyio +async def test_interrupt_exhausted_retries_surface_as_conflict_error(): + """When all retry attempts collide with a unique violation, the loop must + surface ConflictError (HTTP 409) — matching the reject branch — instead of + leaking the raw IntegrityError (HTTP 500). + + Without the post-loop conversion, the last attempt's ``raise`` re-raises + the IntegrityError, giving callers an inconsistent signal depending on + which strategy they picked. The reject path already converts; this test + pins the symmetric behaviour for interrupt/rollback. + """ + import sqlite3 + + class _AlwaysUniqueViolationStore(MemoryRunStore): + """MemoryRunStore whose ``create_run_atomic`` always raises a + real-flavoured unique-violation IntegrityError, simulating a worker + that keeps losing the cross-worker race for the same thread.""" + + def __init__(self): + super().__init__() + self.atomic_call_count = 0 + + async def create_run_atomic(self, *args, **kwargs): + self.atomic_call_count += 1 + err = sqlite3.IntegrityError("UNIQUE constraint failed: runs.uq_runs_thread_active") + err.sqlite_errorcode = sqlite3.SQLITE_CONSTRAINT_UNIQUE + raise err + + store = _AlwaysUniqueViolationStore() + manager = _make_manager(store=store) + + with pytest.raises(ConflictError, match="already has an active run"): + await manager.create_or_reject("thread-1", multitask_strategy="interrupt") + + # Sanity: the loop actually retried 3 times before giving up. + assert store.atomic_call_count == 3 + + +# --------------------------------------------------------------------------- +# create_or_reject — run ownership metadata +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_run_record_stores_owner_and_lease(): + """Newly created runs must carry owner_worker_id and lease_expires_at (when heartbeat is on).""" + store = MemoryRunStore() + manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True)) + record = await manager.create_or_reject("thread-1") + + assert record.owner_worker_id == manager.worker_id + assert isinstance(record.owner_worker_id, str) and len(record.owner_worker_id) > 0 + assert record.lease_expires_at is not None + + # Store row must also carry the fields + stored = await store.get(record.run_id) + assert stored is not None + assert stored["owner_worker_id"] == manager.worker_id + assert stored["lease_expires_at"] is not None + + +@pytest.mark.anyio +async def test_store_row_roundtrips_ownership_fields(): + """Records hydrated from the store must surface ownership fields.""" + store = MemoryRunStore() + manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True)) + record = await manager.create_or_reject("thread-1") + + hydrated = await manager.get(record.run_id) + assert hydrated is not None + assert hydrated.owner_worker_id == manager.worker_id + assert hydrated.lease_expires_at is not None + + +# --------------------------------------------------------------------------- +# reconcile_orphaned_inflight_runs — lease-based +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_reconciliation_claims_expired_lease_runs(): + """A run with an expired lease must be reclaimed as orphaned.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + + # Insert a run with an already-expired lease + expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat() + await store.put( + "expired-run", + thread_id="thread-1", + status="running", + owner_worker_id="worker-dead", + lease_expires_at=expired_lease, + created_at=(datetime.now(UTC) - timedelta(seconds=120)).isoformat(), + ) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="Gateway restarted before this run reached a durable final state.", + ) + + assert len(recovered) == 1 + assert recovered[0].run_id == "expired-run" + assert recovered[0].status == RunStatus.error + + stored = await store.get("expired-run") + assert stored["status"] == "error" + + +@pytest.mark.anyio +async def test_reconciliation_skips_active_lease_runs(): + """A run with a still-valid lease must NOT be reclaimed.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + + # Insert a run with a still-valid lease + valid_lease = (datetime.now(UTC) + timedelta(seconds=60)).isoformat() + await store.put( + "live-run", + thread_id="thread-1", + status="running", + owner_worker_id="worker-alive", + lease_expires_at=valid_lease, + created_at=(datetime.now(UTC) - timedelta(seconds=10)).isoformat(), + ) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="Gateway restarted before this run reached a durable final state.", + ) + + # Live run's lease is still valid — must not be reclaimed + assert all(r.run_id != "live-run" for r in recovered) + + stored = await store.get("live-run") + assert stored["status"] == "running" + + +@pytest.mark.anyio +async def test_reconciliation_claims_null_lease_runs(): + """Pre-ownership rows (NULL lease) must be reclaimed.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + + await store.put( + "legacy-run", + thread_id="thread-1", + status="running", + created_at=(datetime.now(UTC) - timedelta(seconds=120)).isoformat(), + ) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="Gateway restarted before this run reached a durable final state.", + ) + + assert len(recovered) == 1 + assert recovered[0].run_id == "legacy-run" + + +@pytest.mark.anyio +async def test_heartbeat_disabled_crashed_run_reclaimed_immediately(): + """Single-worker regression: when heartbeat is off, a crashed run must be + reclaimed on the next restart without waiting for lease expiry. + + The run is created with lease_expires_at=NULL (no heartbeat => no lease), + so reconciliation treats it as an orphan and reclaims it right away — + preserving the pre-ownership recovery latency. + """ + store = MemoryRunStore() + # Worker A: heartbeat disabled (single-worker default) + manager_a = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=False)) + record = await manager_a.create("thread-1") + await manager_a.set_status(record.run_id, RunStatus.running) + + # Verify the run was stored WITHOUT a lease (heartbeat off) + stored = await store.get(record.run_id) + assert stored is not None + assert stored["lease_expires_at"] is None + + # Simulate crash: drop manager_a's local state, build a fresh manager + # (same store) as if Worker A restarted. + manager_b = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=False)) + + # Reconciliation must reclaim the run IMMEDIATELY — no lease to wait out. + recovered = await manager_b.reconcile_orphaned_inflight_runs( + error="Gateway restarted before this run reached a durable final state.", + ) + + assert len(recovered) == 1 + assert recovered[0].run_id == record.run_id + assert recovered[0].status == RunStatus.error + + +@pytest.mark.anyio +async def test_reconciliation_skips_locally_active_runs(): + """An active local run (owned by this worker) must NOT be reclaimed even with an expired lease.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + + # Create a live local run + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + + # Its lease hasn't expired yet, so this is mostly testing the local-ownership guard + recovered = await manager.reconcile_orphaned_inflight_runs( + error="Gateway restarted before this run reached a durable final state.", + ) + + assert all(r.run_id != record.run_id for r in recovered) + + +@pytest.mark.anyio +async def test_reconciliation_returns_empty_when_no_orphaned_runs(): + """Reconciliation must return empty when there are no orphaned runs.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + + recovered = await manager.reconcile_orphaned_inflight_runs( + error="Gateway restarted before this run reached a durable final state.", + ) + + assert recovered == [] + + +# --------------------------------------------------------------------------- +# Lease heartbeat +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_heartbeat_renews_active_run_leases(): + """Heartbeat must extend the lease on active runs owned by this worker.""" + config = _lease_config(lease_seconds=30, heartbeat_enabled=True) + store = MemoryRunStore() + manager = _make_manager(store=store, run_ownership_config=config) + + record = await manager.create_or_reject("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + + original_lease = record.lease_expires_at + assert original_lease is not None + + # Start heartbeat and let it tick once + await manager.start_heartbeat() + await asyncio.sleep(0.2) # heartbeat interval = 10s, too long; manually renew + + await manager._renew_leases() + await manager.stop_heartbeat() + + assert record.lease_expires_at is not None + # Lease should have been extended + assert record.lease_expires_at >= original_lease + + +@pytest.mark.anyio +async def test_heartbeat_renews_pending_run_before_task_is_spawned(): + """A run sitting in ``pending`` between ``create_run_atomic`` and task + spawn must still have its lease renewed. + + Pre-fix the renewal filter required ``record.task is not None``, so a + pending run with no task yet (the brief window after + ``create_run_atomic`` inserts the row before the worker layer spawns + the agent task) was silently skipped. If that window stretched past + ``lease_seconds`` — e.g. event-loop saturation, slow checkpoint + hydrate — peer reconciliation reclaimed the run as an orphan and + marked it ``error`` even though this worker still intended to run it. + """ + config = _lease_config(lease_seconds=30, heartbeat_enabled=True) + store = MemoryRunStore() + manager = _make_manager(store=store, run_ownership_config=config) + + record = await manager.create_or_reject("thread-1") + assert record.status == RunStatus.pending + # No task has been spawned — this is the regression sentinel. + assert record.task is None + + original_lease = record.lease_expires_at + assert original_lease is not None + + # Force a measurable gap so the renewed lease strictly post-dates the + # original — without this the two timestamps land in the same + # microsecond on fast hosts and the strict comparison fails trivially. + await asyncio.sleep(0.001) + + store.update_lease = AsyncMock(wraps=store.update_lease) + + await manager._renew_leases() + + store.update_lease.assert_awaited_once() + assert record.lease_expires_at is not None + assert record.lease_expires_at > original_lease + + +@pytest.mark.anyio +async def test_heartbeat_skips_runs_not_owned_by_this_worker(): + """Heartbeat must only renew leases for runs owned by this worker.""" + config = _lease_config(lease_seconds=30, heartbeat_enabled=True) + store = MemoryRunStore() + manager = _make_manager(store=store, run_ownership_config=config) + + # Create a run owned by a different worker + old_lease = (datetime.now(UTC) + timedelta(seconds=5)).isoformat() + await store.put( + "other-worker-run", + thread_id="thread-1", + status="running", + owner_worker_id="other-worker", + lease_expires_at=old_lease, + created_at=(datetime.now(UTC) - timedelta(seconds=10)).isoformat(), + ) + + await manager._renew_leases() + + stored = await store.get("other-worker-run") + # Lease should be unchanged (other worker's run) + assert stored["lease_expires_at"] == old_lease + + +@pytest.mark.anyio +async def test_heartbeat_not_started_when_disabled(): + """When heartbeat_enabled is False, start_heartbeat must be a no-op.""" + config = _lease_config(heartbeat_enabled=False) + store = MemoryRunStore() + manager = _make_manager(store=store, run_ownership_config=config) + + assert manager.heartbeat_enabled is False + await manager.start_heartbeat() + assert manager._heartbeat_task is None + assert manager._heartbeat_stop is None + + +# --------------------------------------------------------------------------- +# cancel with cross-worker lease awareness +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_cancel_local_run_succeeds(): + """Cancel must succeed for a locally-owned active run.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.running) + + result = await manager.cancel(record.run_id) + assert result is True + assert record.status == RunStatus.interrupted + + +@pytest.mark.anyio +async def test_cancel_unknown_run_returns_false(): + """Cancel must return False for a run not known to this worker.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + + result = await manager.cancel("nonexistent-run") + assert result is False + + +@pytest.mark.anyio +async def test_cancel_idempotent(): + """Cancel must return True when the run is already interrupted.""" + store = MemoryRunStore() + manager = _make_manager(store=store) + record = await manager.create("thread-1") + await manager.set_status(record.run_id, RunStatus.interrupted) + + result = await manager.cancel(record.run_id) + assert result is True + + +# --------------------------------------------------------------------------- +# GATEWAY_WORKERS=1 backward compatibility +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_single_worker_default_config_behavior_unchanged(): + """With default config (heartbeat_enabled=False), behavior must match pre-ownership code.""" + config = _lease_config(heartbeat_enabled=False) + store = MemoryRunStore() + manager = _make_manager(store=store, run_ownership_config=config) + + # Create runs, cancel, create_or_reject — all must work + r1 = await manager.create("thread-1") + assert r1.owner_worker_id is not None + + r2 = await manager.create_or_reject("thread-2", multitask_strategy="reject") + assert r2.owner_worker_id is not None + + await manager.cancel(r2.run_id) + stored = await store.get(r2.run_id) + assert stored["status"] == "interrupted" + + +@pytest.mark.anyio +async def test_manager_without_run_ownership_config(): + """Manager without run_ownership_config must still work (backward compat).""" + store = MemoryRunStore() + manager = RunManager(store=store) # no run_ownership_config + + record = await manager.create_or_reject("thread-1") + assert record is not None + assert record.owner_worker_id is not None # always set, even without config + + # Heartbeat must be a no-op without config + assert manager.heartbeat_enabled is False + await manager.start_heartbeat() + assert manager._heartbeat_task is None + + +# --------------------------------------------------------------------------- +# worker_id uniqueness +# --------------------------------------------------------------------------- + + +def test_worker_id_is_generated(): + """worker_id must be a non-empty string containing hostname.""" + wid = _generate_worker_id() + assert isinstance(wid, str) + assert len(wid) > 0 + assert ":" in wid + + +def test_two_managers_have_different_default_ids(): + """Two managers without explicit worker_id must get unique ids.""" + m1 = RunManager() + m2 = RunManager() + assert m1.worker_id != m2.worker_id + + +# --------------------------------------------------------------------------- +# Store atomic methods +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_create_run_atomic_reject_prevents_duplicate(): + """store.create_run_atomic with reject must raise ConflictError on duplicate.""" + store = MemoryRunStore() + config = _lease_config() + + store.create_run_atomic = AsyncMock(wraps=store.create_run_atomic) + + await store.create_run_atomic( + run_id="run-1", + thread_id="thread-1", + owner_worker_id="w1", + lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(), + multitask_strategy="reject", + grace_seconds=config.grace_seconds, + ) + + with pytest.raises(ConflictError, match="already has an active run"): + await store.create_run_atomic( + run_id="run-2", + thread_id="thread-1", + owner_worker_id="w2", + lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(), + multitask_strategy="reject", + grace_seconds=config.grace_seconds, + ) + + +@pytest.mark.anyio +async def test_create_run_atomic_interrupt_claims_and_creates(): + """store.create_run_atomic with interrupt must claim old and create new.""" + store = MemoryRunStore() + config = _lease_config() + # Create an active run with an expired lease (simulating a crashed worker) + expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat() + + await store.create_run_atomic( + run_id="run-old", + thread_id="thread-1", + owner_worker_id="w1", + lease_expires_at=expired_lease, + multitask_strategy="reject", + grace_seconds=config.grace_seconds, + ) + + new_row, claimed = await store.create_run_atomic( + run_id="run-new", + thread_id="thread-1", + owner_worker_id="w2", + lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(), + multitask_strategy="interrupt", + grace_seconds=config.grace_seconds, + ) + + assert new_row["run_id"] == "run-new" + assert new_row["status"] == "pending" + assert len(claimed) == 1 + assert claimed[0]["run_id"] == "run-old" + + # Old run must be interrupted in-store + old_row = await store.get("run-old") + assert old_row["status"] == "interrupted" + + +@pytest.mark.anyio +async def test_create_run_atomic_interrupt_rejects_other_worker_valid_lease(): + """Interrupt must raise ConflictError when a valid-lease run is owned by another worker. + + The partial unique index ``uq_runs_thread_active`` would reject the INSERT + anyway; surfacing ConflictError here gives the caller a clean signal + instead of a futile retry loop on IntegrityError. + """ + store = MemoryRunStore() + config = _lease_config(grace_seconds=10) + valid_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat() + + await store.create_run_atomic( + run_id="valid-lease-run", + thread_id="thread-1", + owner_worker_id="other-worker", + lease_expires_at=valid_lease, + multitask_strategy="reject", + grace_seconds=config.grace_seconds, + ) + + with pytest.raises(ConflictError, match="another worker"): + await store.create_run_atomic( + run_id="run-new", + thread_id="thread-1", + owner_worker_id="w2", + lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(), + multitask_strategy="interrupt", + grace_seconds=config.grace_seconds, + ) + + # The valid-lease run must be untouched (transaction rolled back). + old_row = await store.get("valid-lease-run") + assert old_row["status"] == "pending" + assert old_row["owner_worker_id"] == "other-worker" + + +@pytest.mark.anyio +async def test_create_run_atomic_interrupt_allows_self_owned_valid_lease(): + """Interrupt must succeed when the existing valid-lease run is owned by this worker.""" + store = MemoryRunStore() + config = _lease_config(grace_seconds=10) + valid_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat() + + await store.create_run_atomic( + run_id="self-run", + thread_id="thread-1", + owner_worker_id="w1", + lease_expires_at=valid_lease, + multitask_strategy="reject", + grace_seconds=config.grace_seconds, + ) + + new_row, claimed = await store.create_run_atomic( + run_id="run-new", + thread_id="thread-1", + owner_worker_id="w1", # same worker + lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(), + multitask_strategy="interrupt", + grace_seconds=config.grace_seconds, + ) + + assert new_row["run_id"] == "run-new" + assert len(claimed) == 1 + assert claimed[0]["run_id"] == "self-run" + assert claimed[0]["status"] == "interrupted" + + +@pytest.mark.anyio +async def test_create_run_atomic_interrupt_rolls_back_earlier_mutations_on_conflict(): + """Interrupt must not leave earlier candidates interrupted when a later + candidate raises ConflictError. + + Mirrors the SQL store's transactional semantics: the whole interrupt pass + is one transaction, so a raise on any candidate must roll back mutations + already applied to earlier candidates. Without this, the memory store + diverges from SQL (which the production path uses), and the + test_multi_worker_run_ownership.py suite gives false confidence by + passing against memory while SQL would behave differently. + + Setup: expired-lease run (interruptible) inserted FIRST, then a + valid-lease run owned by another worker. Iteration order means the + expired run is mutated before the valid-lease run raises — so a naive + single-pass implementation would leave the expired run interrupted. + """ + store = MemoryRunStore() + config = _lease_config(grace_seconds=10) + expired_lease = (datetime.now(UTC) - timedelta(seconds=60)).isoformat() + valid_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat() + + # Seed both active rows directly via ``put`` (bypassing create_run_atomic's + # reject check, which would refuse the second row). Insert the + # interruptible run first so dict iteration visits it first — that's the + # ordering that exposes the half-interrupted divergence in a naive + # single-pass implementation. + await store.put( + "expired-run", + thread_id="thread-1", + status="pending", + owner_worker_id="old-worker", + lease_expires_at=expired_lease, + ) + await store.put( + "valid-lease-run", + thread_id="thread-1", + status="pending", + owner_worker_id="other-worker", + lease_expires_at=valid_lease, + ) + + with pytest.raises(ConflictError, match="another worker"): + await store.create_run_atomic( + run_id="run-new", + thread_id="thread-1", + owner_worker_id="w1", + lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(), + multitask_strategy="interrupt", + grace_seconds=config.grace_seconds, + ) + + # The expired run must be UNTOUCHED — the interrupt pass must roll back + # on ConflictError, not leave a half-interrupted store. + expired_row = await store.get("expired-run") + assert expired_row["status"] == "pending" + assert expired_row["owner_worker_id"] == "old-worker" + assert expired_row["error"] is None + + # The valid-lease run that caused the conflict is also untouched. + valid_row = await store.get("valid-lease-run") + assert valid_row["status"] == "pending" + assert valid_row["owner_worker_id"] == "other-worker" + + # The new run was never inserted. + assert await store.get("run-new") is None + + +# --------------------------------------------------------------------------- +# update_lease +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_update_lease_renews_row(): + """update_lease must update the lease_expires_at on the stored row.""" + store = MemoryRunStore() + old_lease = (datetime.now(UTC) + timedelta(seconds=5)).isoformat() + await store.put( + "run-1", + thread_id="thread-1", + status="running", + owner_worker_id="w1", + lease_expires_at=old_lease, + ) + + new_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat() + updated = await store.update_lease( + "run-1", + owner_worker_id="w1", + lease_expires_at=new_lease, + ) + assert updated is True + + stored = await store.get("run-1") + assert stored["lease_expires_at"] == new_lease + + +@pytest.mark.anyio +async def test_update_lease_returns_false_for_terminal_run(): + """update_lease must return False when the run is not pending/running.""" + store = MemoryRunStore() + await store.put("run-1", thread_id="thread-1", status="success", owner_worker_id="w1") + + new_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat() + updated = await store.update_lease( + "run-1", + owner_worker_id="w1", + lease_expires_at=new_lease, + ) + assert updated is False + + stored = await store.get("run-1") + assert stored["status"] == "success" + + +@pytest.mark.anyio +async def test_update_lease_returns_false_for_wrong_owner(): + """update_lease must reject renewal when owner_worker_id does not match.""" + store = MemoryRunStore() + old_lease = (datetime.now(UTC) + timedelta(seconds=5)).isoformat() + await store.put( + "run-1", + thread_id="thread-1", + status="running", + owner_worker_id="w1", + lease_expires_at=old_lease, + ) + + new_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat() + updated = await store.update_lease( + "run-1", + owner_worker_id="w2", # different worker + lease_expires_at=new_lease, + ) + assert updated is False + + # The original lease must be untouched + stored = await store.get("run-1") + assert stored["owner_worker_id"] == "w1" + assert stored["lease_expires_at"] == old_lease + + +# --------------------------------------------------------------------------- +# list_inflight_with_expired_lease +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_list_inflight_with_expired_lease_filters_correctly(): + """Only runs with expired or NULL leases must be returned.""" + store = MemoryRunStore() + now = datetime.now(UTC) + grace = 10 + + # Expired lease + expired = (now - timedelta(seconds=60)).isoformat() + await store.put("expired-run", thread_id="t1", status="running", owner_worker_id="w1", lease_expires_at=expired, created_at=expired) + + # Valid lease + valid = (now + timedelta(seconds=60)).isoformat() + await store.put("valid-run", thread_id="t2", status="running", owner_worker_id="w2", lease_expires_at=valid, created_at=valid) + + # NULL lease (legacy) + await store.put("null-lease-run", thread_id="t3", status="running", created_at=(now - timedelta(seconds=30)).isoformat()) + + # Terminal status (should not appear) + await store.put("success-run", thread_id="t4", status="success", created_at=(now - timedelta(seconds=60)).isoformat()) + + results = await store.list_inflight_with_expired_lease(grace_seconds=grace) + + result_ids = {r["run_id"] for r in results} + assert "expired-run" in result_ids + assert "null-lease-run" in result_ids + assert "valid-run" not in result_ids + assert "success-run" not in result_ids + + +# --------------------------------------------------------------------------- +# MemoryRunStore — datetime comparison for created_at filtering +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_list_inflight_with_expired_lease_compares_created_at_as_datetime(): + """``before`` filter must use datetime comparison, not string lexical order. + + ISO-8601 strings compare lexically only when every component is zero-padded + to the same width and the timezone suffix matches. Datetime parsing is + order-safe regardless of format. + """ + store = MemoryRunStore() + now = datetime.now(UTC) + grace = 10 + + # A run created "now" — should be included when before=None (defaults to now). + await store.put("recent-run", thread_id="t1", status="running", created_at=now.isoformat()) + # A run created far in the future — should be excluded by the before filter + # even though the string "2300-01-01..." > "2025-..." lexically. + far_future = "2300-01-01T00:00:00+00:00" + await store.put("future-run", thread_id="t2", status="running", created_at=far_future) + + results = await store.list_inflight_with_expired_lease(before=now.isoformat(), grace_seconds=grace) + result_ids = {r["run_id"] for r in results} + assert "recent-run" in result_ids + assert "future-run" not in result_ids + + +@pytest.mark.anyio +async def test_list_inflight_with_expired_lease_handles_malformed_created_at(): + """Malformed ``created_at`` values must not crash the listing.""" + store = MemoryRunStore() + grace = 10 + + store._runs["bad-run"] = { + "run_id": "bad-run", + "thread_id": "t1", + "status": "running", + "created_at": "not-a-datetime", + } + store._runs["empty-run"] = { + "run_id": "empty-run", + "thread_id": "t2", + "status": "running", + "created_at": "", + } + + results = await store.list_inflight_with_expired_lease(grace_seconds=grace) + # Both should be skipped because their created_at can't be parsed + result_ids = {r["run_id"] for r in results} + assert "bad-run" not in result_ids + assert "empty-run" not in result_ids + + +@pytest.mark.anyio +async def test_list_inflight_with_expired_lease_datetime_aware_naive_handling(): + """Lease comparison must handle aware and naive datetimes. + + ``lease_expires_at`` stored with a trailing ``+00:00`` (aware) and without + (naive) should both be comparable against the aware ``cutoff``. The MemoryRunStore + uses ``datetime.fromisoformat`` which preserves the offset, so both paths + must work. + """ + store = MemoryRunStore() + now = datetime.now(UTC) + grace = 10 + + # Naive datetime (no timezone suffix) — common on SQLite read-back + naive_expired = (now - timedelta(seconds=60)).isoformat() # "2025-01-01T00:00:00" + await store.put("naive-run", thread_id="t1", status="running", lease_expires_at=naive_expired, created_at=naive_expired) + + # Aware datetime (with +00:00) + aware_expired = (now - timedelta(seconds=60)).replace(tzinfo=UTC).isoformat() # "2025-01-01T00:00:00+00:00" + await store.put("aware-run", thread_id="t2", status="running", lease_expires_at=aware_expired, created_at=aware_expired) + + results = await store.list_inflight_with_expired_lease(grace_seconds=grace) + result_ids = {r["run_id"] for r in results} + # Both expired, both should be returned + assert "naive-run" in result_ids + assert "aware-run" in result_ids + + +@pytest.mark.anyio +async def test_list_inflight_with_expired_lease_null_lease_always_reclaimed(): + """NULL lease rows are always reclaimed regardless of created_at value.""" + store = MemoryRunStore() + grace = 10 + + # NULL lease is the single-worker mode default — every inflight row + # must be returned so reconciliation can reclaim it. + await store.put("null-run", thread_id="t1", status="running", created_at=datetime.now(UTC).isoformat()) + + results = await store.list_inflight_with_expired_lease(grace_seconds=grace) + result_ids = {r["run_id"] for r in results} + assert "null-run" in result_ids diff --git a/backend/tests/test_owner_isolation.py b/backend/tests/test_owner_isolation.py index 33d21f3e3..ac190bbdf 100644 --- a/backend/tests/test_owner_isolation.py +++ b/backend/tests/test_owner_isolation.py @@ -164,8 +164,8 @@ async def test_runs_cross_user_isolation(tmp_path): repo = RunRepository(get_session_factory()) with _as_user(USER_A): - await repo.put("run-a1", thread_id="t-alpha") - await repo.put("run-a2", thread_id="t-alpha") + await repo.put("run-a1", thread_id="t-alpha", status="success") + await repo.put("run-a2", thread_id="t-alpha", status="pending") with _as_user(USER_B): await repo.put("run-b1", thread_id="t-beta") diff --git a/backend/tests/test_persistence_bootstrap.py b/backend/tests/test_persistence_bootstrap.py index 439702510..0a9d8b0f6 100644 --- a/backend/tests/test_persistence_bootstrap.py +++ b/backend/tests/test_persistence_bootstrap.py @@ -47,7 +47,7 @@ from deerflow.persistence.migrations._helpers import _normalize_default asyncio_test = pytest.mark.asyncio -HEAD = "0003_scheduled_tasks" +HEAD = "0004_run_ownership" BASELINE = "0001_baseline" @@ -74,6 +74,11 @@ async def _runs_column_meta(engine, column_name: str) -> dict: raise AssertionError(f"column {column_name!r} not found in runs") +async def _runs_index_names(engine) -> set[str]: + async with engine.connect() as conn: + return await conn.run_sync(lambda c: {ix["name"] for ix in sa.inspect(c).get_indexes("runs")}) + + async def _alembic_version(engine) -> str | None: async with engine.connect() as conn: row = await conn.execute(sa.text("SELECT version_num FROM alembic_version")) @@ -142,6 +147,12 @@ async def test_empty_branch_creates_all_and_stamps_head(tmp_path: Path) -> None: assert required in tables, f"missing table: {required}" assert "token_usage_by_model" in await _runs_columns(engine) assert await _alembic_version(engine) == HEAD + # The partial unique index on (thread_id WHERE status IN pending/running) + # must exist on a fresh DB because the empty-branch stamps head without + # running migrations, so the index has to come from ``Base.metadata``. + indexes = await _runs_index_names(engine) + assert "uq_runs_thread_active" in indexes, indexes + assert "ix_runs_lease" in indexes, indexes finally: await engine.dispose() diff --git a/backend/tests/test_persistence_bootstrap_concurrency.py b/backend/tests/test_persistence_bootstrap_concurrency.py index 23a00e06c..de41df28f 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 = "0003_scheduled_tasks" +HEAD = "0004_run_ownership" 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 a9d8c47c2..16683db58 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] == "0003_scheduled_tasks" + assert version_row[0] == "0004_run_ownership" # 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] == "0003_scheduled_tasks" + assert version_row[0] == "0004_run_ownership" finally: await close_engine() diff --git a/backend/tests/test_run_manager.py b/backend/tests/test_run_manager.py index 820dcf384..b1c5dfd2f 100644 --- a/backend/tests/test_run_manager.py +++ b/backend/tests/test_run_manager.py @@ -643,7 +643,7 @@ async def test_create_or_reject_does_not_interrupt_old_run_when_new_run_store_wr manager = RunManager(store=store) old = await manager.create("thread-1") await manager.set_status(old.run_id, RunStatus.running) - store.put = AsyncMock(side_effect=RuntimeError("db down")) + store.create_run_atomic = AsyncMock(side_effect=RuntimeError("db down")) with pytest.raises(RuntimeError, match="db down"): await manager.create_or_reject("thread-1", multitask_strategy="interrupt") @@ -664,10 +664,10 @@ async def test_create_or_reject_does_not_interrupt_old_run_when_new_run_store_wr old = await manager.create("thread-1") await manager.set_status(old.run_id, RunStatus.running) - async def cancelled_put(run_id, **kwargs): + async def cancelled_create(run_id, **kwargs): raise asyncio.CancelledError - store.put = cancelled_put + store.create_run_atomic = cancelled_create with pytest.raises(asyncio.CancelledError): await manager.create_or_reject("thread-1", multitask_strategy="interrupt") @@ -881,11 +881,14 @@ async def test_list_by_thread_falls_back_to_store_with_user_filter(): class _FailingPutRunStore(MemoryRunStore): - """Memory run store whose every ``put`` fails (non-retryably).""" + """Memory run store whose every ``put`` and ``create_run_atomic`` fails (non-retryably).""" async def put(self, run_id, **kwargs): raise ValueError("simulated persist failure") + async def create_run_atomic(self, run_id, **kwargs): + raise ValueError("simulated persist failure") + @pytest.mark.anyio async def test_thread_index_scopes_runs_per_thread(manager: RunManager): diff --git a/backend/tests/test_run_repository.py b/backend/tests/test_run_repository.py index 7b44cbb93..f9c975a72 100644 --- a/backend/tests/test_run_repository.py +++ b/backend/tests/test_run_repository.py @@ -8,6 +8,7 @@ from sqlalchemy.dialects import postgresql from deerflow.persistence.run import RunRepository from deerflow.runtime import RunManager, RunStatus +from deerflow.runtime.runs.manager import ConflictError from deerflow.runtime.runs.store.base import RunStore @@ -56,6 +57,15 @@ class _CustomRunStoreWithoutProgress(RunStore): async def aggregate_tokens_by_thread(self, *args, **kwargs): return {} + async def update_lease(self, *args, **kwargs): + return True + + async def list_inflight_with_expired_lease(self, *args, **kwargs): + return [] + + async def create_run_atomic(self, *args, **kwargs): + return {}, [] + @pytest.mark.anyio async def test_update_run_progress_defaults_to_noop_for_custom_store(): @@ -125,9 +135,9 @@ class TestRunRepository: @pytest.mark.anyio async def test_list_by_thread(self, tmp_path): repo = await _make_repo(tmp_path) - await repo.put("r1", thread_id="t1") - await repo.put("r2", thread_id="t1") - await repo.put("r3", thread_id="t2") + await repo.put("r1", thread_id="t1", status="success") + await repo.put("r2", thread_id="t1", status="pending") + await repo.put("r3", thread_id="t2", status="pending") rows = await repo.list_by_thread("t1") assert len(rows) == 2 assert all(r["thread_id"] == "t1" for r in rows) @@ -136,8 +146,8 @@ class TestRunRepository: @pytest.mark.anyio async def test_list_by_thread_owner_filter(self, tmp_path): repo = await _make_repo(tmp_path) - await repo.put("r1", thread_id="t1", user_id="alice") - await repo.put("r2", thread_id="t1", user_id="bob") + await repo.put("r1", thread_id="t1", user_id="alice", status="success") + await repo.put("r2", thread_id="t1", user_id="bob", status="pending") rows = await repo.list_by_thread("t1", user_id="alice") assert len(rows) == 1 assert rows[0]["user_id"] == "alice" @@ -161,8 +171,8 @@ class TestRunRepository: async def test_list_pending(self, tmp_path): repo = await _make_repo(tmp_path) await repo.put("r1", thread_id="t1", status="pending") - await repo.put("r2", thread_id="t1", status="running") - await repo.put("r3", thread_id="t2", status="pending") + await repo.put("r2", thread_id="t2", status="running") + await repo.put("r3", thread_id="t3", status="pending") pending = await repo.list_pending() assert len(pending) == 2 assert all(r["status"] == "pending" for r in pending) @@ -171,10 +181,13 @@ class TestRunRepository: @pytest.mark.anyio async def test_list_inflight_returns_pending_and_running_before_cutoff(self, tmp_path): repo = await _make_repo(tmp_path) + # Each thread can hold at most one pending/running row (partial unique + # index ``uq_runs_thread_active``), so spread the inflight rows across + # distinct threads to exercise the before-cutoff filter. await repo.put("pending-old", thread_id="t1", status="pending", created_at="2026-01-01T00:00:00+00:00") - await repo.put("running-old", thread_id="t1", status="running", created_at="2026-01-01T00:00:01+00:00") - await repo.put("success-old", thread_id="t1", status="success", created_at="2026-01-01T00:00:02+00:00") - await repo.put("pending-new", thread_id="t1", status="pending", created_at="2026-01-01T00:00:03+00:00") + await repo.put("running-old", thread_id="t2", status="running", created_at="2026-01-01T00:00:01+00:00") + await repo.put("success-old", thread_id="t3", status="success", created_at="2026-01-01T00:00:02+00:00") + await repo.put("pending-new", thread_id="t4", status="pending", created_at="2026-01-01T00:00:03+00:00") inflight = await repo.list_inflight(before="2026-01-01T00:00:02+00:00") @@ -394,8 +407,8 @@ class TestRunRepository: async def test_list_by_thread_ordered_desc(self, tmp_path): """list_by_thread returns newest first.""" repo = await _make_repo(tmp_path) - await repo.put("r1", thread_id="t1", created_at="2024-01-01T00:00:00+00:00") - await repo.put("r2", thread_id="t1", created_at="2024-01-02T00:00:00+00:00") + await repo.put("r1", thread_id="t1", status="success", created_at="2024-01-01T00:00:00+00:00") + await repo.put("r2", thread_id="t1", status="pending", created_at="2024-01-02T00:00:00+00:00") rows = await repo.list_by_thread("t1") assert rows[0]["run_id"] == "r2" assert rows[1]["run_id"] == "r1" @@ -404,8 +417,11 @@ class TestRunRepository: @pytest.mark.anyio async def test_list_by_thread_limit(self, tmp_path): repo = await _make_repo(tmp_path) - for i in range(5): - await repo.put(f"r{i}", thread_id="t1") + # Only one row can be pending/running per thread; mark earlier ones + # terminal so the partial unique index still holds. + for i in range(4): + await repo.put(f"r{i}", thread_id="t1", status="success") + await repo.put("r4", thread_id="t1", status="pending") rows = await repo.list_by_thread("t1", limit=2) assert len(rows) == 2 await _cleanup() @@ -413,8 +429,8 @@ class TestRunRepository: @pytest.mark.anyio async def test_owner_none_returns_all(self, tmp_path): repo = await _make_repo(tmp_path) - await repo.put("r1", thread_id="t1", user_id="alice") - await repo.put("r2", thread_id="t1", user_id="bob") + await repo.put("r1", thread_id="t1", user_id="alice", status="success") + await repo.put("r2", thread_id="t1", user_id="bob", status="pending") rows = await repo.list_by_thread("t1", user_id=None) assert len(rows) == 2 await _cleanup() @@ -428,21 +444,21 @@ class TestRunRepository: await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path)) repo = RunRepository(get_session_factory()) - await repo.put("run-1", thread_id="thread-1", model_name="gpt-4o") + await repo.put("run-1", thread_id="thread-1", model_name="gpt-4o", status="success") row = await repo.get("run-1") assert row is not None assert row["model_name"] == "gpt-4o" long_name = "a" * 200 - await repo.put("run-2", thread_id="thread-1", model_name=long_name) + await repo.put("run-2", thread_id="thread-1", model_name=long_name, status="success") row2 = await repo.get("run-2") assert row2["model_name"] == "a" * 128 - await repo.put("run-3", thread_id="thread-1", model_name=123) + await repo.put("run-3", thread_id="thread-1", model_name=123, status="success") row3 = await repo.get("run-3") assert row3["model_name"] == "123" - await repo.put("run-4", thread_id="thread-1", model_name=None) + await repo.put("run-4", thread_id="thread-1", model_name=None, status="pending") row4 = await repo.get("run-4") assert row4["model_name"] is None @@ -625,3 +641,180 @@ class TestRunRepository: row = await repo.get(record.run_id) assert row["model_name"] == "model-2" await _cleanup() + + @pytest.mark.anyio + async def test_create_run_atomic_reject_propagates_conflict_on_unique_violation(self, tmp_path): + """reject path against a real SQLite-backed store must surface as ConflictError, not raw IntegrityError. + + The partial unique index ``uq_runs_thread_active`` is created by + ``Base.metadata.create_all`` on SQLite too. Every other atomic-create + test in the suite uses ``MemoryRunStore``, which raises ConflictError + directly and never exercises the manager's + ``_is_unique_violation``-based conversion. This test is the load-bearing + coverage for that branch on a real DB: pre-insert an active run on + thread T, then attempt a reject-strategy create for the same thread, + and assert ConflictError (HTTP 409) — not a leaking IntegrityError + (HTTP 500). + """ + from datetime import UTC, datetime, timedelta + + from deerflow.config.run_ownership_config import RunOwnershipConfig + + repo = await _make_repo(tmp_path) + manager = RunManager( + store=repo, + run_ownership_config=RunOwnershipConfig( + lease_seconds=30, + grace_seconds=10, + heartbeat_enabled=False, + ), + ) + + # Pre-insert an active run on thread T directly through the store so + # the partial unique index has something to enforce on the second insert. + lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat() + await repo.create_run_atomic( + "run-A", + thread_id="thread-T", + owner_worker_id="worker-A", + lease_expires_at=lease, + multitask_strategy="reject", + created_at=datetime.now(UTC).isoformat(), + ) + + # Second reject-strategy create against the same thread must convert the + # underlying IntegrityError into ConflictError via ``_is_unique_violation``. + with pytest.raises(ConflictError, match="already has an active run"): + await manager.create_or_reject( + "thread-T", + multitask_strategy="reject", + ) + + await _cleanup() + + @pytest.mark.anyio + async def test_is_unique_violation_detects_real_sqlite_integrity_error(self, tmp_path): + """``_is_unique_violation`` must return True for a real SQLite IntegrityError. + + SQLite raises ``UNIQUE constraint failed: runs.uq_runs_thread_active`` + which contains "unique" but neither "violat" nor "duplicate" — the + previous substring-only heuristic returned False on SQLite, leaking the + raw IntegrityError. This test triggers a real violation against the + partial unique index and feeds the resulting SQLAlchemy IntegrityError + (with the wrapped sqlite3.IntegrityError on ``.orig``) through the + detector to assert True. + """ + import sqlite3 + + from sqlalchemy.exc import IntegrityError + + from deerflow.runtime.runs.manager import _is_unique_violation + + repo = await _make_repo(tmp_path) + + # First insert succeeds; second collides on the partial unique index. + await repo.put("first", thread_id="thread-T", status="pending") + with pytest.raises(IntegrityError) as exc_info: + await repo.put("second", thread_id="thread-T", status="pending") + + # The wrapped driver exception must be a sqlite3 IntegrityError carrying + # SQLITE_CONSTRAINT_UNIQUE. Walk the chain so we assert on the actual + # driver-level signal, not the SQLAlchemy wrapper. + driver = exc_info.value.orig + assert isinstance(driver, sqlite3.IntegrityError) + assert driver.sqlite_errorcode == sqlite3.SQLITE_CONSTRAINT_UNIQUE + + # The detector must return True regardless of message phrasing. + assert _is_unique_violation(exc_info.value) is True + + await _cleanup() + + @pytest.mark.anyio + async def test_is_unique_violation_does_not_misclassify_application_exception(self): + """Message fallbacks must not fire on non-IntegrityError exceptions. + + A ``ValueError`` / ``RuntimeError`` whose ``str()`` happens to + contain ``"duplicate key"`` or ``"unique" + "violat"`` substrings + must NOT be classified as a unique violation — that would silently + mask real application bugs as HTTP 409 conflicts instead of 500. + Pre-fix the substring-only fallback fired regardless of exception + type. The fix gates the fallback on + ``isinstance(current, (SAIntegrityError, sqlite3.IntegrityError))``. + """ + from deerflow.runtime.runs.manager import _is_unique_violation + + assert _is_unique_violation(ValueError("duplicate key in input data: 'email'")) is False + assert _is_unique_violation(RuntimeError("unique violat detected in config")) is False + assert _is_unique_violation(Exception("unique constraint failed (in a unit test mock)")) is False + + @pytest.mark.anyio + async def test_is_unique_violation_detects_psycopg3_sqlstate(self): + """psycopg3 exposes the error code via ``sqlstate``, not ``pgcode``. + + On Postgres (the only supported multi-worker backend), psycopg3's + ``sqlstate=23505`` must be detected as a unique violation without + falling through to the message-substring fallback. + """ + from sqlalchemy.exc import IntegrityError as SAIntegrityError + + from deerflow.runtime.runs.manager import _is_unique_violation + + # Simulate psycopg3's sqlstate attribute on a wrapped IntegrityError + dbapi_err = Exception() + dbapi_err.sqlstate = "23505" # psycopg3 uses sqlstate + + sa_err = SAIntegrityError( + "duplicate key value violates unique constraint", + params=None, + orig=dbapi_err, + ) + + assert _is_unique_violation(sa_err) is True + + @pytest.mark.anyio + async def test_create_run_atomic_interrupt_tolerates_tz_naive_lease_on_sqlite(self, tmp_path): + """Interrupt path must not raise TypeError comparing naive vs aware datetimes. + + SQLite drops tzinfo on read despite ``DateTime(timezone=True)`` (see + the comment in ``RunRepository._row_to_dict``). The interrupt branch + of ``create_run_atomic`` compares ``row.lease_expires_at`` against + the aware ``cutoff = datetime.now(UTC) - ...`` in Python. Under + default config (heartbeat disabled) leases are always NULL so the + ``is not None`` check short-circuits, but there is no guard against + ``heartbeat_enabled=true`` on SQLite — a naive lease would raise + ``TypeError: can't compare offset-naive and offset-aware datetimes`` + and surface as an opaque 500. + + Pre-fix this test fails with TypeError; post-fix it raises + ConflictError (the live other-worker run blocks the interrupt). + """ + from datetime import UTC, datetime, timedelta + + repo = await _make_repo(tmp_path) + + # Seed an active run owned by another worker with a still-valid lease. + # The lease value is stored as ISO; SQLite reads it back as a tz-naive + # datetime — exactly the shape that triggered the bug. + valid_lease = (datetime.now(UTC) + timedelta(seconds=30)).isoformat() + await repo.create_run_atomic( + "valid-lease-run", + thread_id="thread-T", + owner_worker_id="other-worker", + lease_expires_at=valid_lease, + multitask_strategy="reject", + created_at=datetime.now(UTC).isoformat(), + ) + + # The interrupt path must surface a clean ConflictError, not a + # TypeError from the naive-vs-aware comparison. + with pytest.raises(ConflictError, match="another worker"): + await repo.create_run_atomic( + "run-new", + thread_id="thread-T", + owner_worker_id="w1", + lease_expires_at=(datetime.now(UTC) + timedelta(seconds=30)).isoformat(), + multitask_strategy="interrupt", + created_at=datetime.now(UTC).isoformat(), + ) + + await _cleanup() diff --git a/config.example.yaml b/config.example.yaml index b5409f1ec..b8316e6ad 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -1596,6 +1596,30 @@ scheduler: max_concurrent_runs: 3 min_once_delay_seconds: 60 +# ============================================================================ +# Run Ownership Configuration +# ============================================================================ +# Controls cross-process run ownership for multi-worker deployments. +# When GATEWAY_WORKERS > 1, each worker claims runs with a lease; the heartbeat +# renews leases, and reconciliation recovers orphaned runs from crashed workers. +# +# CLOCK-SYNC REQUIREMENT (multi-worker only): reconciliation compares another +# worker's UTC lease timestamp against this worker's datetime.now(UTC). Worker +# clocks MUST be synced (NTP / chrony / systemd-timesyncd — default on K8s and +# cloud VMs) within a few seconds. grace_seconds is the skew budget; worst case +# (owning worker's heartbeat just about to fire), a peer whose clock is more +# than grace_seconds ahead can mis-reclaim a still-live run as an orphan. Raise +# grace_seconds if your environment cannot keep clocks within a few seconds; +# the trade-off is longer recovery latency for genuinely dead workers +# (lease_seconds + grace_seconds from last heartbeat to reclaim). + +run_ownership: + lease_seconds: 30 # Seconds before a run lease expires if not renewed. + # Heartbeat renews every lease_seconds / 3. + grace_seconds: 10 # Extra seconds past expiry before reclaiming an orphaned run. + # Also the cross-worker clock-skew budget — see note above. + heartbeat_enabled: false # Set to true for GATEWAY_WORKERS > 1 + # ============================================================================ # Stream Bridge Configuration # ============================================================================