From 828363705a20dbc2d49f5dc9388c02fef1574c17 Mon Sep 17 00:00:00 2001 From: rain02333z-spec Date: Fri, 14 Aug 2026 23:50:40 +0800 Subject: [PATCH] fix(scheduler): support safe multi-instance scheduler recovery (#4713) * fix(scheduler): reject unsafe multi-worker startup * fix(scheduler): support safe multi-instance recovery * fix(scheduler): make multi-instance recovery lease-safe * fix(scheduler): address multi-instance review feedback * docs(scheduler): add multi-instance upgrade notes --------- Co-authored-by: rain02333z-spec <225106191+rain02333z-spec@users.noreply.github.com> --- README.md | 8 + backend/AGENTS.md | 1 + backend/app/gateway/app.py | 2 + backend/app/gateway/deps.py | 46 +- backend/app/scheduler/service.py | 133 ++++- backend/docs/CONFIGURATION.md | 8 +- .../deerflow/config/reload_boundary.py | 3 +- .../deerflow/config/scheduler_config.py | 1 + .../persistence/scheduled_task_runs/sql.py | 85 ++- .../persistence/scheduled_tasks/sql.py | 219 +++++++- backend/tests/test_gateway_run_recovery.py | 30 + .../tests/test_multi_worker_postgres_gate.py | 103 ++++ backend/tests/test_scheduled_task_models.py | 1 + backend/tests/test_scheduled_task_postgres.py | 148 +++++ .../tests/test_scheduled_task_repository.py | 518 +++++++++++++++++- backend/tests/test_scheduled_task_service.py | 116 +++- config.example.yaml | 12 +- deploy/helm/deer-flow/README.md | 11 +- deploy/helm/deer-flow/values.yaml | 6 +- 19 files changed, 1359 insertions(+), 92 deletions(-) create mode 100644 backend/tests/test_scheduled_task_postgres.py diff --git a/README.md b/README.md index 3c3bc08e4..8444d5ec5 100644 --- a/README.md +++ b/README.md @@ -1297,6 +1297,14 @@ Current MVP limits: Enable background polling with `config.yaml -> scheduler.enabled`. Manual trigger uses the same scheduled-task resource and execution path. +The background scheduler is single-instance by default. For a multi-pod deployment, set `scheduler.multi_instance: true` and use shared Postgres, `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; startup and periodic recovery then preserve live peer runs, atomically take over only expired leases, and fence stale post-launch writes. `max_concurrent_runs` is a shared global cap across Pods, including short-lived dispatch reservations. Without those settings, enable the scheduler on exactly one Gateway pod. These scheduler fields are startup-only; restart all Gateway Pods together when changing them. + +### Upgrade Notes + +- Before upgrading a deployment with `GATEWAY_WORKERS > 1` and `scheduler.enabled: true`, either keep the scheduler on exactly one Gateway worker or configure `scheduler.multi_instance: true` with shared Postgres, `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`. The upgraded Gateway rejects the unsafe combination at startup instead of starting silently. +- In multi-instance mode, `scheduler.max_concurrent_runs` is a cluster-wide cap, not a per-Pod cap. It includes active scheduled runs and short-lived dispatch reservations, so capacity does not multiply with the number of replicas. +- `scheduler.multi_instance` and the related scheduler, ownership, and run-event settings are startup-only. Apply changes with a coordinated restart of all Gateway Pods; changing the ConfigMap alone does not activate multi-instance recovery. + ## Terminal Workbench (TUI) `deerflow` is a terminal-native workbench for people who live in the shell. It runs **embedded** over `DeerFlowClient` — no Gateway, frontend, nginx, or Docker required — while honoring the same `config.yaml`, checkpointer, skills, memory, MCP, and sandbox settings as the rest of DeerFlow. diff --git a/backend/AGENTS.md b/backend/AGENTS.md index f9dd77655..f1a5ae044 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -18,6 +18,7 @@ DeerFlow is a LangGraph-based AI super agent system with a full-stack architectu - With `stream_subgraphs`, subgraph frames keep their namespace in the SSE event name (`values|`, LangGraph Platform style) instead of impersonating root frames — a delegated subagent inherits the parent checkpoint namespace, so publishing its `values` snapshot as bare `values` replaces the whole thread view in SDK clients (#4399). Root-only consumers (file-tool chunk batcher, subagent event persistence, LLM error-fallback detection) ignore namespaced frames. The web frontend does not request subgraph streaming; subtask progress rides root-namespace `task_*` custom events. - Background subagent identity is deliberately split: the provider `tool_call_id` remains the correlation key for `ToolMessage`, `task_*` SSE events, persisted lifecycle events, frontend cards, and the public `ExtensionData.scope_id` contract (stored as `SubagentResult.external_task_id`), while `SubagentExecutor.execute_async()` generates a full server-side `execution_id` for `SubagentResult.task_id`, the process-wide registry, polling, cancellation, timeout handling, and cleanup. Provider IDs are not globally unique across parent runs, so they must never become registry ownership keys; scheduler closures retain their own `SubagentResult` rather than resolving ownership again through the mutable registry. Terminal subagent token usage travels in the current run's `ToolMessage.additional_kwargs` and is attributed from message state, never through a process-global provider-ID cache. - Scheduled-task executions must reuse that same Gateway run lifecycle. The scheduler may decide *when* work runs, but it must dispatch through the existing run path rather than introducing a parallel execution stack. +- The background scheduler is single-instance by default. `scheduler.multi_instance=true` opts into lease-aware recovery across Gateway instances and requires shared Postgres, `run_ownership.heartbeat_enabled=true`, and `run_events.backend=db`; otherwise startup rejects the configuration. Live scheduled runs are preserved when a peer starts; expired leases are atomically taken over, stale post-launch writes are fenced by the dispatch lease owner, and the Postgres advisory-locked budget makes `max_concurrent_runs` a shared global cap (including pre-launch reservations). - Long-running MCP work uses a separate durable task runtime rather than keeping remote task IDs or status polling inside the Agent loop. `McpTaskService` claims due rows with leases, resolves a protocol-specific `McpTaskDriver`, and writes normalized snapshots back to `mcp_tasks`; expired leases are the restart-recovery mechanism, and a result returned after expiry must be discarded even when the owner token still matches. The database is the source of truth. `ThreadState` may receive only a bounded projection in later integration work, never the sole recoverable copy. - Scheduled-task dispatch enforces "at most one active run per task when `overlap_policy=skip`" at the DB layer via the partial unique index `uq_scheduled_task_run_active` (`scheduled_task_runs.task_id WHERE status IN ('queued','running')`). `ScheduledTaskService.dispatch_task`'s `has_active_runs` check is a non-atomic fast path (its own session, separated from the `create()` insert by `await` points), so two concurrent dispatches — a manual `POST /scheduled-tasks/{id}/trigger` racing the poller, a double-click, or a client retry — can both pass it; the index is the atomic arbiter, and the losing `create` surfaces as `ActiveScheduledRunConflict` (translated from `IntegrityError` in the repository) and collapses to the same outcome as the fast path (manual → 409 conflict, scheduled → a `"skipped"` tombstone). The scheduled-skip tombstone is created directly as terminal `"skipped"` (not a transient `"queued"`) so it never occupies the active slot the pre-existing run still holds. Sibling of the `runs` table's `uq_runs_thread_active` (PR #4003), which keys on `thread_id` and so does not cover the default `fresh_thread_per_run` context where every dispatch gets a new thread. Index is status-only, not `overlap_policy`-conditional (the policy is fixed to `"skip"` in the MVP). diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index 5cba7f55d..9cc6f8b81 100644 --- a/backend/app/gateway/app.py +++ b/backend/app/gateway/app.py @@ -322,6 +322,8 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: poll_interval_seconds=startup_config.scheduler.poll_interval_seconds, lease_seconds=startup_config.scheduler.lease_seconds, max_concurrent_runs=startup_config.scheduler.max_concurrent_runs, + multi_instance=startup_config.scheduler.multi_instance, + run_lease_grace_seconds=startup_config.run_ownership.grace_seconds, ) app.state.scheduled_task_service = scheduled_task_service if startup_config.scheduler.enabled: diff --git a/backend/app/gateway/deps.py b/backend/app/gateway/deps.py index 36b468fc9..3a17275f9 100644 --- a/backend/app/gateway/deps.py +++ b/backend/app/gateway/deps.py @@ -56,18 +56,21 @@ def _browser_tools_enabled_in_config(config: AppConfig) -> bool: def _enforce_postgres_for_multi_worker(config: AppConfig) -> None: - """Refuse unsafe multi-worker configurations before persistence starts. + """Refuse unsafe multi-process configurations before persistence starts. - Four checks (all must pass for multi-worker): + Multi-instance scheduler recovery also needs the durable run ownership + contract even when each Pod runs a single Gateway worker. - 1. Process-local browser sessions must be disabled. Browser tools keep + 1. The background scheduler must be disabled for ordinary multi-worker + mode. ``scheduler.multi_instance`` opts into the lease-aware path. + 2. Process-local browser sessions must be disabled. Browser tools keep Chromium and Playwright objects in one worker's memory, while ordinary uvicorn dispatch provides no thread-id affinity. - 2. The DB backend must be Postgres — SQLite write-locks cannot support + 3. The DB backend must be Postgres — SQLite write-locks cannot support concurrent multi-process access. - 3. ``run_events.backend`` must be ``db``. Memory and JSONL stores are + 4. ``run_events.backend`` must be ``db``. Memory and JSONL stores are process-local, so workers cannot enforce a shared singleton receipt. - 4. ``run_ownership.heartbeat_enabled`` must be True — without heartbeat, + 5. ``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. @@ -81,17 +84,33 @@ def _enforce_postgres_for_multi_worker(config: AppConfig) -> None: except (TypeError, ValueError): workers = 1 + scheduler = getattr(config, "scheduler", None) + multi_instance_requested = bool(getattr(scheduler, "multi_instance", False)) + multi_instance_scheduler = bool(getattr(scheduler, "enabled", False) and multi_instance_requested) + + backend = getattr(config.database, "backend", None) + run_events_backend = getattr(getattr(config, "run_events", None), "backend", None) + run_ownership = getattr(config, "run_ownership", None) + + if multi_instance_requested and backend != "postgres": + raise SystemExit(f"scheduler.multi_instance=true requires database.backend='postgres'. database.backend is '{backend}'. Set scheduler.multi_instance=false or configure Postgres.") + if multi_instance_requested and run_events_backend != "db": + raise SystemExit(f"scheduler.multi_instance=true requires run_events.backend='db'. run_events.backend is '{run_events_backend}'. Set scheduler.multi_instance=false or configure run_events.backend: db.") + if multi_instance_requested and (run_ownership is None or not run_ownership.heartbeat_enabled): + raise SystemExit("scheduler.multi_instance=true requires run_ownership.heartbeat_enabled=true so peer runs retain a valid lease. Set scheduler.multi_instance=false or enable run ownership heartbeats.") + if workers <= 1: return + if config.scheduler.enabled and not multi_instance_scheduler: + raise SystemExit(f"GATEWAY_WORKERS={workers} cannot run with scheduler.enabled=true because each worker starts its own scheduler. Set GATEWAY_WORKERS=1, scheduler.multi_instance=true, or scheduler.enabled=false.") + if _browser_tools_enabled_in_config(config): raise SystemExit(browser_multi_worker_error(workers)) - backend = getattr(config.database, "backend", 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_events_backend = getattr(getattr(config, "run_events", None), "backend", None) if run_events_backend != "db": raise SystemExit( f"GATEWAY_WORKERS={workers} requires run_events.backend='db', but run_events.backend is '{run_events_backend}'. " @@ -99,7 +118,6 @@ def _enforce_postgres_for_multi_worker(config: AppConfig) -> None: "Set GATEWAY_WORKERS=1 or configure run_events.backend: db." ) - 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. " @@ -481,9 +499,15 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen ) from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository + app.state.scheduled_task_repo = ScheduledTaskRepository( + sf, + run_repository=app.state.run_store, + ) + app.state.scheduled_task_run_repo = ScheduledTaskRunRepository( + sf, + run_repository=app.state.run_store, + ) app.state.mcp_task_repo = McpTaskRepository(sf) - app.state.scheduled_task_repo = ScheduledTaskRepository(sf) - app.state.scheduled_task_run_repo = ScheduledTaskRunRepository(sf) else: app.state.mcp_task_repo = None app.state.scheduled_task_repo = None diff --git a/backend/app/scheduler/service.py b/backend/app/scheduler/service.py index 2a6574f03..8d8ce71f3 100644 --- a/backend/app/scheduler/service.py +++ b/backend/app/scheduler/service.py @@ -20,6 +20,8 @@ logger = logging.getLogger(__name__) # byte-identical outcomes for the same "task already has an active run" condition. _ACTIVE_RUN_CONFLICT_ERROR = "task already has an active run" _SKIP_ACTIVE_RUN_ERROR = "skipped: a previous run of this task is still active" +_RESTART_RECOVERY_ERROR = "interrupted: gateway restarted before the run reached a terminal state" +_LEASE_RECOVERY_ERROR = "interrupted: the owning gateway stopped renewing its run lease" class ScheduledTaskService: @@ -32,6 +34,8 @@ class ScheduledTaskService: poll_interval_seconds: int, lease_seconds: int, max_concurrent_runs: int, + multi_instance: bool = False, + run_lease_grace_seconds: int = 10, ) -> None: self._task_repo = task_repo self._task_run_repo = task_run_repo @@ -39,24 +43,40 @@ class ScheduledTaskService: self._poll_interval_seconds = poll_interval_seconds self._lease_seconds = lease_seconds self._max_concurrent_runs = max_concurrent_runs + self._multi_instance = multi_instance + self._run_lease_grace_seconds = run_lease_grace_seconds self._lease_owner = f"{socket.gethostname()}:{uuid.uuid4().hex}" self._task: asyncio.Task | None = None self._stop = asyncio.Event() + self._skip_next_lease_reconciliation = False async def run_once(self, *, now: datetime) -> None: - # ``max_concurrent_runs`` is a global cap on active scheduled runs, not - # just a per-poll claim batch: long runs accumulate across poll cycles, - # so each cycle only claims into the remaining budget. - active = await self._task_run_repo.count_active_runs() - budget = self._max_concurrent_runs - active - if budget <= 0: - return - claimed = await self._task_repo.claim_due_tasks( - now=now, - lease_owner=self._lease_owner, - lease_seconds=self._lease_seconds, - limit=budget, - ) + if self._multi_instance: + if self._skip_next_lease_reconciliation: + self._skip_next_lease_reconciliation = False + else: + await self._reconcile_active_state(now=now) + claimed = await self._task_repo.claim_due_tasks( + now=now, + lease_owner=self._lease_owner, + lease_seconds=self._lease_seconds, + limit=self._max_concurrent_runs, + global_max_concurrent_runs=self._max_concurrent_runs, + ) + else: + # In single-instance mode the count and claim do not need a shared + # database lock. Multi-instance mode performs both inside the + # repository's short Postgres advisory-lock transaction above. + active = await self._task_run_repo.count_active_runs() + budget = self._max_concurrent_runs - active + if budget <= 0: + return + claimed = await self._task_repo.claim_due_tasks( + now=now, + lease_owner=self._lease_owner, + lease_seconds=self._lease_seconds, + limit=budget, + ) for task in claimed: await self.dispatch_task(task, now=now, trigger="scheduled") @@ -105,6 +125,7 @@ class ScheduledTaskService: now: datetime, trigger: str, ) -> dict[str, Any]: + expected_lease_owner = self._lease_owner if trigger == "scheduled" else None execution_thread_id = task.get("thread_id") if task.get("context_mode") == "fresh_thread_per_run" or execution_thread_id is None: execution_thread_id = str(uuid.uuid4()) @@ -133,6 +154,7 @@ class ScheduledTaskService: last_thread_id=execution_thread_id, last_error=str(exc), increment_run_count=False, + expected_lease_owner=expected_lease_owner, ) return { "outcome": "failed", @@ -161,6 +183,17 @@ class ScheduledTaskService: return self._active_run_conflict_result(execution_thread_id) return await self._record_scheduled_skip(task, thread_id=execution_thread_id, now=now, trigger=trigger) + if self._multi_instance and trigger == "manual": + task = await self._task_repo.claim_dispatch_lease( + task["id"], + lease_owner=self._lease_owner, + now=now, + lease_seconds=self._lease_seconds, + ) + if task is None: + return self._active_run_conflict_result(execution_thread_id) + expected_lease_owner = self._lease_owner + task_run_id = f"task-run-{uuid.uuid4().hex}" try: await self._task_run_repo.create( @@ -236,6 +269,7 @@ class ScheduledTaskService: # Same race as the run-row write above: a fast-failing run's # completion hook may have already finalized a `once` task. protect_terminal=True, + expected_lease_owner=expected_lease_owner, ) return { "outcome": "launched", @@ -256,6 +290,7 @@ class ScheduledTaskService: thread_id=execution_thread_id, now=now, error=str(exc), + trigger=trigger, ) next_at = next_run_at( @@ -309,6 +344,7 @@ class ScheduledTaskService: last_error=None, increment_run_count=True, protect_terminal=True, + expected_lease_owner=expected_lease_owner, ) except Exception: logger.exception( @@ -343,6 +379,7 @@ class ScheduledTaskService: last_thread_id=execution_thread_id, last_error=str(exc), increment_run_count=False, + expected_lease_owner=expected_lease_owner, ) return { "outcome": "conflict" if self._is_overlap_conflict(exc) else "failed", @@ -391,7 +428,14 @@ class ScheduledTaskService: trigger=trigger, status="skipped", ) - return await self._finalize_skip(task, task_run_id=task_run_id, thread_id=thread_id, now=now, error=_SKIP_ACTIVE_RUN_ERROR) + return await self._finalize_skip( + task, + task_run_id=task_run_id, + thread_id=thread_id, + now=now, + error=_SKIP_ACTIVE_RUN_ERROR, + trigger=trigger, + ) async def _finalize_skip( self, @@ -401,6 +445,7 @@ class ScheduledTaskService: thread_id: str, now: datetime, error: str, + trigger: str, ) -> dict[str, Any]: next_at = next_run_at( task["schedule_type"], @@ -424,6 +469,7 @@ class ScheduledTaskService: last_thread_id=task.get("last_thread_id"), last_error=error if task["schedule_type"] == "once" else None, increment_run_count=False, + expected_lease_owner=self._lease_owner if trigger == "scheduled" else None, ) return { "outcome": "skipped", @@ -487,25 +533,52 @@ class ScheduledTaskService: async def start(self) -> None: if self._task is not None: return - restart_error = "interrupted: gateway restarted before the run reached a terminal state" - try: - stale = await self._task_run_repo.mark_stale_active_runs(error=restart_error) - if stale: - logger.warning("Marked %d stale scheduled task run(s) as interrupted after restart", stale) - except Exception: - logger.exception("Failed to sweep stale scheduled task runs at startup") - try: - # The run rows above are only half the story: a launched `once` - # task is parked in "running" until the (now dead) completion hook - # would have finalized it, so reconcile the parent rows too. - stuck = await self._task_repo.cancel_stuck_once_tasks(error=restart_error) - if stuck: - logger.warning("Cancelled %d stuck once task(s) after restart", stuck) - except Exception: - logger.exception("Failed to reconcile stuck once tasks at startup") + restart_error = _RESTART_RECOVERY_ERROR + if self._multi_instance: + await self._reconcile_active_state(now=datetime.now(UTC)) + self._skip_next_lease_reconciliation = True + else: + try: + stale = await self._task_run_repo.mark_stale_active_runs(error=restart_error) + if stale: + logger.warning("Marked %d stale scheduled task run(s) as interrupted after restart", stale) + except Exception: + logger.exception("Failed to sweep stale scheduled task runs at startup") + try: + # The run rows above are only half the story: a launched `once` + # task is parked in "running" until the (now dead) completion hook + # would have finalized it, so reconcile the parent rows too. + stuck = await self._task_repo.cancel_stuck_once_tasks(error=restart_error) + if stuck: + logger.warning("Cancelled %d stuck once task(s) after restart", stuck) + except Exception: + logger.exception("Failed to reconcile stuck once tasks at startup") self._stop.clear() self._task = asyncio.create_task(self._run_loop()) + async def _reconcile_active_state(self, *, now: datetime) -> None: + error = _LEASE_RECOVERY_ERROR + try: + stale = await self._task_run_repo.reconcile_active_runs( + error=error, + now=now, + lease_grace_seconds=self._run_lease_grace_seconds, + ) + if stale: + logger.warning("Marked %d stale scheduled task run(s) as interrupted after lease reconciliation", stale) + except Exception: + logger.exception("Failed to reconcile scheduled task runs with leases") + try: + stuck = await self._task_repo.reconcile_stuck_once_tasks( + error=error, + now=now, + lease_grace_seconds=self._run_lease_grace_seconds, + ) + if stuck: + logger.warning("Cancelled %d stuck once task(s) after lease reconciliation", stuck) + except Exception: + logger.exception("Failed to reconcile once tasks with leases") + async def stop(self) -> None: if self._task is None: return diff --git a/backend/docs/CONFIGURATION.md b/backend/docs/CONFIGURATION.md index d128ba043..4ad0f0268 100644 --- a/backend/docs/CONFIGURATION.md +++ b/backend/docs/CONFIGURATION.md @@ -237,6 +237,7 @@ The scheduled-task MVP adds a scheduler section to `config.yaml`: ```yaml scheduler: enabled: false + multi_instance: false poll_interval_seconds: 5 lease_seconds: 120 max_concurrent_runs: 3 @@ -246,8 +247,13 @@ scheduler: Notes: - `enabled: false` keeps background polling off by default. -- `max_concurrent_runs` is a global cap on active scheduled runs (queued/running run rows); each poll cycle claims only into the remaining budget, so long runs accumulating across cycles cannot exceed it. +- `multi_instance: true` opts into lease-aware scheduler recovery across Gateway instances. It requires Postgres, `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`; otherwise startup fails fast. Leave it false for the default single-instance scheduler. +- `max_concurrent_runs` is a shared global cap in multi-instance mode. It counts active `queued`/`running` scheduled-run rows plus valid pre-launch dispatch leases, and Postgres serializes the budget read with due-task claims so long runs or concurrent Pods cannot exceed it. +- Multi-instance reconciliation uses the run ownership lease: a live peer run is preserved, an expired lease is atomically taken over before its scheduled row is interrupted, and a stale Pod cannot overwrite a newer Pod's parent-task bookkeeping. - All scheduler fields are restart-required; edits need a Gateway restart. +- **Upgrade note:** before upgrading a deployment with `GATEWAY_WORKERS > 1` and `scheduler.enabled: true`, either run the scheduler on exactly one Gateway worker or enable `scheduler.multi_instance: true` with shared Postgres, `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`. The startup gate now rejects the unsafe combination instead of allowing it to start silently. +- **Upgrade note:** in multi-instance mode, `max_concurrent_runs` is cluster-wide rather than per Pod and includes active scheduled runs plus dispatch reservations. Plan capacity accordingly; it does not multiply with the replica count. +- **Upgrade note:** `scheduler.multi_instance` and its related scheduler, ownership, and run-event settings are startup-only. Restart all Gateway Pods together after changing them; a ConfigMap update without a coordinated restart leaves the running service on its previous mode. - Multi-worker deployments (`GATEWAY_WORKERS > 1`) must use the Postgres database backend, enable run ownership heartbeats, and set `run_events.backend: db`. SQLite silently ignores row-level locks, while memory and JSONL run-event stores are process-local and cannot enforce singleton delivery receipts across workers; startup rejects these combinations. The process-local agentic browser tool group is incompatible with multiple Gateway workers; keep `GATEWAY_WORKERS=1` while `browser_navigate` is enabled. Browser control also requires the backend `browser` extra (`cd backend && uv sync --extra browser && uv run playwright install chromium`); startup detects enabled browser config and fails fast when Playwright is missing, and `/api/features` reports `browser_control.enabled=false` until the runtime is available. - The MVP supports thread reuse and fresh-thread-per-run execution modes. - The MVP supports only `once` and `cron`. diff --git a/backend/packages/harness/deerflow/config/reload_boundary.py b/backend/packages/harness/deerflow/config/reload_boundary.py index 89247abdf..e4ed7995c 100644 --- a/backend/packages/harness/deerflow/config/reload_boundary.py +++ b/backend/packages/harness/deerflow/config/reload_boundary.py @@ -68,7 +68,8 @@ STARTUP_ONLY_FIELDS: dict[str, str] = { ), "scheduler": ( "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." + "max_concurrent_runs, and multi_instance are captured into the service instance and the background poller task is not rebuilt on config.yaml edits. " + "Changing multi-instance recovery prerequisites or lease behavior requires restarting every Gateway Pod together." ), "mcp_tasks": ( "McpTaskService is constructed and started once during Gateway lifespan startup; enabled, poll_interval_seconds, lease_seconds, " diff --git a/backend/packages/harness/deerflow/config/scheduler_config.py b/backend/packages/harness/deerflow/config/scheduler_config.py index dfc0bb207..ac51453ab 100644 --- a/backend/packages/harness/deerflow/config/scheduler_config.py +++ b/backend/packages/harness/deerflow/config/scheduler_config.py @@ -3,6 +3,7 @@ from pydantic import BaseModel, Field class SchedulerConfig(BaseModel): enabled: bool = Field(default=False) + multi_instance: bool = Field(default=False) poll_interval_seconds: int = Field(default=5, ge=1, le=300) lease_seconds: int = Field(default=120, ge=5, le=3600) max_concurrent_runs: int = Field(default=3, ge=1, le=32) diff --git a/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py b/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py index afedb4b4f..1b04f4ac3 100644 --- a/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py +++ b/backend/packages/harness/deerflow/persistence/scheduled_task_runs/sql.py @@ -1,19 +1,30 @@ from __future__ import annotations -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from typing import Any from sqlalchemy import func, select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from deerflow.persistence.run import RunRepository +from deerflow.persistence.run.model import RunRow from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow +from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow from deerflow.utils.time import coerce_iso TERMINAL_RUN_STATUSES: frozenset[str] = frozenset({"success", "failed", "skipped", "interrupted"}) ACTIVE_RUN_STATUSES: tuple[str, ...] = ("queued", "running") +def _lease_is_alive(lease_expires_at: datetime | None, *, now: datetime, grace_seconds: int) -> bool: + if lease_expires_at is None: + return False + if lease_expires_at.tzinfo is None: + lease_expires_at = lease_expires_at.replace(tzinfo=UTC) + return lease_expires_at >= now - timedelta(seconds=grace_seconds) + + class ActiveScheduledRunConflict(Exception): """A concurrent dispatch already holds the task's single active-run slot. @@ -35,8 +46,14 @@ class ActiveScheduledRunConflict(Exception): class ScheduledTaskRunRepository: - def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None: + def __init__( + self, + session_factory: async_sessionmaker[AsyncSession], + *, + run_repository: RunRepository | None = None, + ) -> None: self._sf = session_factory + self._run_repository = run_repository or RunRepository(session_factory) @staticmethod def _row_to_dict(row: ScheduledTaskRunRow) -> dict[str, Any]: @@ -169,3 +186,67 @@ class ScheduledTaskRunRepository: row.finished_at = now await session.commit() return len(rows) + + async def reconcile_active_runs( + self, + *, + error: str, + now: datetime, + lease_grace_seconds: int = 10, + ) -> int: + """Reconcile only rows whose underlying owner is no longer live. + + ``RunManager`` owns the durable run lease. A scheduled row with a live + underlying run, or a queued row whose parent task still has a dispatch + lease, belongs to another process and must survive this startup. + """ + async with self._sf() as session: + result = await session.execute(select(ScheduledTaskRunRow.id).where(ScheduledTaskRunRow.status.in_(ACTIVE_RUN_STATUSES))) + row_ids = list(result.scalars()) + stale = 0 + for row_id in row_ids: + row = await session.get(ScheduledTaskRunRow, row_id, with_for_update=True) + if row is None or row.status not in ACTIVE_RUN_STATUSES: + continue + task = await session.get(ScheduledTaskRow, row.task_id, with_for_update=True) + candidate = await self._find_underlying_run(session, row, task) + if candidate is not None and candidate.status in {"pending", "running"}: + if _lease_is_alive(candidate.lease_expires_at, now=now, grace_seconds=lease_grace_seconds): + continue + # Run takeover commits in its own short transaction. If this + # outer commit fails, the next poll finishes scheduled-row + # bookkeeping while the run remains safely terminal. + claimed = await self._run_repository.claim_for_takeover( + candidate.run_id, + grace_seconds=lease_grace_seconds, + error=error, + stop_reason="scheduled_task_orphan_recovered", + ) + if not claimed: + refreshed = await self._run_repository.get(candidate.run_id, user_id=None) + if refreshed is not None and refreshed.get("status") in {"pending", "running"}: + continue + if row.run_id is None and task is not None and _lease_is_alive(task.lease_expires_at, now=now, grace_seconds=0): + continue + row.status = "interrupted" + row.error = error + row.finished_at = now + stale += 1 + await session.commit() + return stale + + @staticmethod + async def _find_underlying_run(session: AsyncSession, row: ScheduledTaskRunRow, task: ScheduledTaskRow | None) -> RunRow | None: + run_ids = [candidate for candidate in (row.run_id, task.last_run_id if task is not None else None) if candidate] + for run_id in dict.fromkeys(run_ids): + candidate = await session.get(RunRow, run_id) + if candidate is None: + continue + linked_task_run_id = (candidate.metadata_json or {}).get("scheduled_task_run_id") + # A stale parent ``last_run_id`` may point at a previous occurrence. + # Let the current scheduled-run metadata lookup recover the live row. + if linked_task_run_id is None or linked_task_run_id == row.id: + return candidate + + result = await session.execute(select(RunRow).where(RunRow.metadata_json["scheduled_task_run_id"].as_string() == row.id).order_by(RunRow.created_at.desc()).limit(1)) + return result.scalars().first() diff --git a/backend/packages/harness/deerflow/persistence/scheduled_tasks/sql.py b/backend/packages/harness/deerflow/persistence/scheduled_tasks/sql.py index 7f05f12bd..46d7ac0e6 100644 --- a/backend/packages/harness/deerflow/persistence/scheduled_tasks/sql.py +++ b/backend/packages/harness/deerflow/persistence/scheduled_tasks/sql.py @@ -1,20 +1,41 @@ from __future__ import annotations +import logging from datetime import UTC, datetime, timedelta from typing import Any -from sqlalchemy import and_, or_, select +from sqlalchemy import and_, exists, func, or_, select, text from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from deerflow.persistence.run import RunRepository +from deerflow.persistence.run.model import RunRow +from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow from deerflow.utils.time import coerce_iso +logger = logging.getLogger(__name__) + TERMINAL_TASK_STATUSES: frozenset[str] = frozenset({"completed", "failed", "cancelled"}) +_SCHEDULER_BUDGET_LOCK_KEY = 4694001 + + +def _lease_is_alive(lease_expires_at: datetime | None, *, now: datetime, grace_seconds: int = 0) -> bool: + if lease_expires_at is None: + return False + if lease_expires_at.tzinfo is None: + lease_expires_at = lease_expires_at.replace(tzinfo=UTC) + return lease_expires_at >= now - timedelta(seconds=grace_seconds) class ScheduledTaskRepository: - def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None: + def __init__( + self, + session_factory: async_sessionmaker[AsyncSession], + *, + run_repository: RunRepository | None = None, + ) -> None: self._sf = session_factory + self._run_repository = run_repository or RunRepository(session_factory) @staticmethod def _row_to_dict(row: ScheduledTaskRow) -> dict[str, Any]: @@ -115,36 +136,60 @@ class ScheduledTaskRepository: lease_owner: str, lease_seconds: int, limit: int, + global_max_concurrent_runs: int | None = None, ) -> list[dict[str, Any]]: lease_expires_at = now + timedelta(seconds=lease_seconds) - stmt = ( - select(ScheduledTaskRow) - .where( - ScheduledTaskRow.next_run_at.is_not(None), - ScheduledTaskRow.next_run_at <= now, - or_( - and_( - ScheduledTaskRow.status == "enabled", - or_( - ScheduledTaskRow.lease_expires_at.is_(None), + async with self._sf() as session: + if global_max_concurrent_runs is not None: + if session.get_bind().dialect.name == "postgresql": + await session.execute(text("SELECT pg_advisory_xact_lock(:lock_key)"), {"lock_key": _SCHEDULER_BUDGET_LOCK_KEY}) + active_runs = await session.scalar(select(func.count()).select_from(ScheduledTaskRunRow).where(ScheduledTaskRunRow.status.in_(("queued", "running")))) + active_run_for_task = exists( + select(ScheduledTaskRunRow.id).where( + ScheduledTaskRunRow.task_id == ScheduledTaskRow.id, + ScheduledTaskRunRow.status.in_(("queued", "running")), + ) + ) + dispatch_reservations = await session.scalar( + select(func.count()) + .select_from(ScheduledTaskRow) + .where( + ScheduledTaskRow.lease_owner.is_not(None), + ScheduledTaskRow.lease_expires_at >= now, + ~active_run_for_task, + ) + ) + active = int(active_runs or 0) + int(dispatch_reservations or 0) + limit = min(limit, max(0, global_max_concurrent_runs - active)) + if limit <= 0: + return [] + stmt = ( + select(ScheduledTaskRow) + .where( + ScheduledTaskRow.next_run_at.is_not(None), + ScheduledTaskRow.next_run_at <= now, + or_( + and_( + ScheduledTaskRow.status == "enabled", + or_( + ScheduledTaskRow.lease_expires_at.is_(None), + ScheduledTaskRow.lease_expires_at < now, + ), + ), + # A task stuck in "running" with an expired lease means the + # claiming process died between claim and dispatch; it must + # stay reclaimable or the task is dead forever. + and_( + ScheduledTaskRow.status == "running", + ScheduledTaskRow.lease_expires_at.is_not(None), ScheduledTaskRow.lease_expires_at < now, ), ), - # A task stuck in "running" with an expired lease means the - # claiming process died between claim and dispatch; it must - # stay reclaimable or the task is dead forever. - and_( - ScheduledTaskRow.status == "running", - ScheduledTaskRow.lease_expires_at.is_not(None), - ScheduledTaskRow.lease_expires_at < now, - ), - ), + ) + .order_by(ScheduledTaskRow.next_run_at.asc(), ScheduledTaskRow.id.asc()) + .limit(limit) + .with_for_update(skip_locked=True) ) - .order_by(ScheduledTaskRow.next_run_at.asc(), ScheduledTaskRow.id.asc()) - .limit(limit) - .with_for_update(skip_locked=True) - ) - async with self._sf() as session: result = await session.execute(stmt) rows = list(result.scalars()) for row in rows: @@ -167,11 +212,21 @@ class ScheduledTaskRepository: last_error: str | None, increment_run_count: bool, protect_terminal: bool = False, - ) -> None: + expected_lease_owner: str | None = None, + ) -> bool: async with self._sf() as session: - row = await session.get(ScheduledTaskRow, task_id) + row = await session.get(ScheduledTaskRow, task_id, with_for_update=True) if row is None: - return + return False + if expected_lease_owner is not None and row.lease_owner != expected_lease_owner: + logger.warning( + "Fenced stale scheduled-task update for task %s: expected lease owner %s, current owner %s", + task_id, + expected_lease_owner, + row.lease_owner, + ) + await session.rollback() + return False if protect_terminal and row.status in TERMINAL_TASK_STATUSES: # A fast-failing run can reach handle_run_completion (which # finalizes a `once` task) before this launch-path write @@ -191,6 +246,38 @@ class ScheduledTaskRepository: row.lease_expires_at = None row.updated_at = datetime.now(UTC) await session.commit() + return True + + async def claim_dispatch_lease( + self, + task_id: str, + *, + lease_owner: str, + now: datetime, + lease_seconds: int, + ) -> dict[str, Any] | None: + """Reserve the short pre-launch window for a manual dispatch.""" + stmt = ( + select(ScheduledTaskRow) + .where( + ScheduledTaskRow.id == task_id, + or_( + ScheduledTaskRow.lease_expires_at.is_(None), + ScheduledTaskRow.lease_expires_at < now, + ), + ) + .with_for_update(skip_locked=True) + ) + async with self._sf() as session: + row = (await session.execute(stmt)).scalars().first() + if row is None: + return None + row.lease_owner = lease_owner + row.lease_expires_at = now + timedelta(seconds=lease_seconds) + row.updated_at = datetime.now(UTC) + await session.commit() + await session.refresh(row) + return self._row_to_dict(row) async def list_by_user_and_thread(self, user_id: str, thread_id: str) -> list[dict[str, Any]]: stmt = ( @@ -230,3 +317,77 @@ class ScheduledTaskRepository: row.updated_at = now await session.commit() return len(rows) + + async def reconcile_stuck_once_tasks( + self, + *, + error: str, + now: datetime, + lease_grace_seconds: int = 10, + ) -> int: + """Cancel once tasks only after their underlying run is no longer live.""" + async with self._sf() as session: + result = await session.execute( + select(ScheduledTaskRow.id).where( + ScheduledTaskRow.schedule_type == "once", + ScheduledTaskRow.status == "running", + ) + ) + task_ids = list(result.scalars()) + cancelled = 0 + for task_id in task_ids: + task = await session.get(ScheduledTaskRow, task_id, with_for_update=True) + if task is None or task.status != "running": + continue + if _lease_is_alive(task.lease_expires_at, now=now, grace_seconds=0): + continue + run_result = await session.execute( + select(ScheduledTaskRunRow) + .where( + ScheduledTaskRunRow.task_id == task.id, + ScheduledTaskRunRow.status.in_(("queued", "running")), + ) + .order_by(ScheduledTaskRunRow.created_at.desc()) + .limit(1) + ) + task_run = run_result.scalars().first() + candidate = await self._find_underlying_run(session, task_run, task) + if candidate is not None and candidate.status in {"pending", "running"}: + if _lease_is_alive(candidate.lease_expires_at, now=now, grace_seconds=lease_grace_seconds): + continue + # Run takeover commits in its own short transaction. If this + # outer commit fails, the next poll finishes task bookkeeping + # while the underlying run remains safely terminal. + claimed = await self._run_repository.claim_for_takeover( + candidate.run_id, + grace_seconds=lease_grace_seconds, + error=error, + stop_reason="scheduled_task_orphan_recovered", + ) + if not claimed: + refreshed = await self._run_repository.get(candidate.run_id, user_id=None) + if refreshed is not None and refreshed.get("status") in {"pending", "running"}: + continue + task.status = "cancelled" + task.last_error = error + task.updated_at = datetime.now(UTC) + cancelled += 1 + await session.commit() + return cancelled + + @staticmethod + async def _find_underlying_run(session: AsyncSession, task_run: ScheduledTaskRunRow | None, task: ScheduledTaskRow) -> RunRow | None: + run_ids = [candidate for candidate in (task_run.run_id if task_run is not None else None, task.last_run_id) if candidate] + for run_id in dict.fromkeys(run_ids): + candidate = await session.get(RunRow, run_id) + if candidate is None: + continue + linked_task_run_id = (candidate.metadata_json or {}).get("scheduled_task_run_id") + if task_run is None or linked_task_run_id is None or linked_task_run_id == task_run.id: + return candidate + + metadata_filter = RunRow.metadata_json["scheduled_task_id"].as_string() == task.id + if task_run is not None: + metadata_filter = RunRow.metadata_json["scheduled_task_run_id"].as_string() == task_run.id + result = await session.execute(select(RunRow).where(metadata_filter).order_by(RunRow.created_at.desc()).limit(1)) + return result.scalars().first() diff --git a/backend/tests/test_gateway_run_recovery.py b/backend/tests/test_gateway_run_recovery.py index 9f6b05c7e..e793ff21b 100644 --- a/backend/tests/test_gateway_run_recovery.py +++ b/backend/tests/test_gateway_run_recovery.py @@ -261,6 +261,36 @@ async def test_sqlite_runtime_reconciles_orphaned_runs_on_startup(monkeypatch): assert stream_bridge.cleanup_calls == [("run-1", 60.0)] +@pytest.mark.anyio +async def test_sql_runtime_shares_run_repository_with_scheduler(monkeypatch): + app = FastAPI() + config = SimpleNamespace( + database=SimpleNamespace(backend="sqlite", checkpoint_channel_mode="full", checkpoint_delta=SimpleNamespace(snapshot_frequency=10)), + run_events=SimpleNamespace(backend="memory"), + stream_bridge=SimpleNamespace(recovered_stream_cleanup_delay_seconds=60.0), + ) + session_factory = object() + _FakeRunManager.instances.clear() + _FakeRunManager.recovered_runs = [] + + async def noop(*_args, **_kwargs): + return None + + monkeypatch.setattr(engine_module, "init_engine_from_config", noop) + monkeypatch.setattr(engine_module, "get_session_factory", lambda: session_factory) + monkeypatch.setattr(engine_module, "close_engine", noop) + monkeypatch.setattr(runtime_module, "make_stream_bridge", lambda _config: _fake_context(_FakeStreamBridge())) + monkeypatch.setattr(checkpointer_module, "make_checkpointer", lambda _config: _fake_context(object())) + monkeypatch.setattr(runtime_module, "make_store", lambda _config: _fake_context(object())) + monkeypatch.setattr(thread_meta_module, "make_thread_store", lambda _sf, _store: _FakeThreadStore()) + monkeypatch.setattr(event_store_module, "make_run_event_store", lambda _config: object()) + monkeypatch.setattr(gateway_deps, "RunManager", _FakeRunManager) + + async with gateway_deps.langgraph_runtime(app, config): + assert app.state.scheduled_task_repo._run_repository is app.state.run_store + assert app.state.scheduled_task_run_repo._run_repository is app.state.run_store + + @pytest.mark.anyio async def test_sqlite_runtime_does_not_mark_thread_error_when_newer_run_is_success(monkeypatch): """Startup recovery should not let an old orphaned run overwrite a newer terminal thread state.""" diff --git a/backend/tests/test_multi_worker_postgres_gate.py b/backend/tests/test_multi_worker_postgres_gate.py index b142fd925..b4b1f1bca 100644 --- a/backend/tests/test_multi_worker_postgres_gate.py +++ b/backend/tests/test_multi_worker_postgres_gate.py @@ -30,6 +30,8 @@ def _config_with_backend( heartbeat_enabled: bool | None = None, browser_enabled: bool = False, run_events_backend: str = "db", + scheduler_enabled: bool = False, + scheduler_multi_instance: bool = False, ) -> SimpleNamespace: run_ownership = RunOwnershipConfig(heartbeat_enabled=heartbeat_enabled) if heartbeat_enabled is not None else None tools = [SimpleNamespace(name="browser_navigate")] if browser_enabled else [] @@ -37,6 +39,7 @@ def _config_with_backend( database=DatabaseConfig(backend=backend), run_ownership=run_ownership, run_events=SimpleNamespace(backend=run_events_backend), + scheduler=SimpleNamespace(enabled=scheduler_enabled, multi_instance=scheduler_multi_instance), tools=tools, ) @@ -65,6 +68,106 @@ def test_gate_allows_multi_worker_with_postgres_and_heartbeat(monkeypatch): _enforce_postgres_for_multi_worker(_config_with_backend("postgres", heartbeat_enabled=True)) +def test_gate_rejects_multi_worker_with_scheduler_enabled(monkeypatch): + monkeypatch.setenv("GATEWAY_WORKERS", "2") + with pytest.raises(SystemExit) as exc_info: + _enforce_postgres_for_multi_worker( + _config_with_backend( + "postgres", + heartbeat_enabled=True, + scheduler_enabled=True, + ) + ) + msg = str(exc_info.value) + assert "scheduler.multi_instance=true" in msg + assert "GATEWAY_WORKERS=1" in msg + assert "scheduler.enabled=false" in msg + + +def test_gate_allows_single_worker_with_scheduler_enabled(monkeypatch): + monkeypatch.setenv("GATEWAY_WORKERS", "1") + _enforce_postgres_for_multi_worker( + _config_with_backend("sqlite", scheduler_enabled=True), + ) + + +def test_gate_allows_multi_instance_scheduler_with_single_worker(monkeypatch): + monkeypatch.setenv("GATEWAY_WORKERS", "1") + _enforce_postgres_for_multi_worker( + _config_with_backend( + "postgres", + heartbeat_enabled=True, + scheduler_enabled=True, + scheduler_multi_instance=True, + ) + ) + + +def test_gate_allows_multi_instance_scheduler_with_multiple_workers(monkeypatch): + monkeypatch.setenv("GATEWAY_WORKERS", "2") + _enforce_postgres_for_multi_worker( + _config_with_backend( + "postgres", + heartbeat_enabled=True, + scheduler_enabled=True, + scheduler_multi_instance=True, + ) + ) + + +def test_gate_rejects_multi_instance_scheduler_without_postgres(monkeypatch): + monkeypatch.setenv("GATEWAY_WORKERS", "1") + with pytest.raises(SystemExit, match="database.backend='postgres'"): + _enforce_postgres_for_multi_worker( + _config_with_backend( + "sqlite", + heartbeat_enabled=True, + scheduler_enabled=True, + scheduler_multi_instance=True, + ) + ) + + +def test_gate_rejects_unsafe_multi_instance_config_even_when_scheduler_disabled(monkeypatch): + monkeypatch.setenv("GATEWAY_WORKERS", "1") + with pytest.raises(SystemExit, match="database.backend='postgres'"): + _enforce_postgres_for_multi_worker( + _config_with_backend( + "sqlite", + heartbeat_enabled=True, + scheduler_enabled=False, + scheduler_multi_instance=True, + ) + ) + + +def test_gate_rejects_multi_instance_scheduler_without_heartbeat(monkeypatch): + monkeypatch.setenv("GATEWAY_WORKERS", "1") + with pytest.raises(SystemExit, match="heartbeat_enabled=true"): + _enforce_postgres_for_multi_worker( + _config_with_backend( + "postgres", + heartbeat_enabled=False, + scheduler_enabled=True, + scheduler_multi_instance=True, + ) + ) + + +def test_gate_rejects_multi_instance_scheduler_with_process_local_events(monkeypatch): + monkeypatch.setenv("GATEWAY_WORKERS", "1") + with pytest.raises(SystemExit, match="run_events.backend='db'"): + _enforce_postgres_for_multi_worker( + _config_with_backend( + "postgres", + heartbeat_enabled=True, + run_events_backend="memory", + scheduler_enabled=True, + scheduler_multi_instance=True, + ) + ) + + @pytest.mark.parametrize("run_events_backend", ["memory", "jsonl"]) def test_gate_rejects_process_local_run_events_with_multi_worker(monkeypatch, run_events_backend): monkeypatch.setenv("GATEWAY_WORKERS", "2") diff --git a/backend/tests/test_scheduled_task_models.py b/backend/tests/test_scheduled_task_models.py index c4f7b1f4b..f8a0d30aa 100644 --- a/backend/tests/test_scheduled_task_models.py +++ b/backend/tests/test_scheduled_task_models.py @@ -10,6 +10,7 @@ def test_app_config_exposes_scheduler_section(): } ) assert config.scheduler.enabled is False + assert config.scheduler.multi_instance is False assert config.scheduler.poll_interval_seconds == 5 assert config.scheduler.lease_seconds == 120 diff --git a/backend/tests/test_scheduled_task_postgres.py b/backend/tests/test_scheduled_task_postgres.py new file mode 100644 index 000000000..62676b17f --- /dev/null +++ b/backend/tests/test_scheduled_task_postgres.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import asyncio +import os +import uuid +from datetime import UTC, datetime, timedelta +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +import pytest +import pytest_asyncio +from sqlalchemy import text + +from deerflow.config.database_config import DatabaseConfig +from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config +from deerflow.persistence.run import RunRepository +from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository +from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository + +POSTGRES_URL = os.environ.get("TEST_POSTGRES_URI") + +pytestmark = pytest.mark.skipif( + not POSTGRES_URL, + reason="requires TEST_POSTGRES_URI (real Postgres for scheduler concurrency)", +) + + +def _postgres_url(url: str) -> str: + parts = urlsplit(url) + query = urlencode([(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key not in {"sslmode", "channel_binding"}]) + return urlunsplit(parts._replace(query=query)) + + +@pytest_asyncio.fixture() +async def postgres_repositories(): + assert POSTGRES_URL is not None + schema = f"scheduler_{uuid.uuid4().hex}" + await init_engine_from_config( + DatabaseConfig( + backend="postgres", + postgres_url=_postgres_url(POSTGRES_URL), + postgres_schema=schema, + ) + ) + sf = get_session_factory() + assert sf is not None + run_repo = RunRepository(sf) + task_repo = ScheduledTaskRepository(sf, run_repository=run_repo) + task_run_repo = ScheduledTaskRunRepository(sf, run_repository=run_repo) + try: + yield task_repo, task_run_repo, run_repo + finally: + engine = get_engine() + async with engine.begin() as connection: + await connection.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')) + await close_engine() + + +async def _create_cron_task(task_repo: ScheduledTaskRepository, task_id: str, *, next_run_at: datetime | None) -> None: + await task_repo.create( + task_id=task_id, + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title=task_id, + prompt="p", + schedule_type="cron", + schedule_spec={"cron": "* * * * *"}, + timezone="UTC", + next_run_at=next_run_at, + ) + + +@pytest.mark.asyncio +async def test_postgres_global_budget_serializes_cross_pod_claims(postgres_repositories): + task_repo, _task_run_repo, _run_repo = postgres_repositories + now = datetime.now(UTC) + await _create_cron_task(task_repo, "task-a", next_run_at=now) + await _create_cron_task(task_repo, "task-b", next_run_at=now) + + claims = await asyncio.gather( + task_repo.claim_due_tasks( + now=now, + lease_owner="pod-a", + lease_seconds=60, + limit=1, + global_max_concurrent_runs=1, + ), + task_repo.claim_due_tasks( + now=now, + lease_owner="pod-b", + lease_seconds=60, + limit=1, + global_max_concurrent_runs=1, + ), + ) + + assert sum(len(claimed) for claimed in claims) == 1 + + +@pytest.mark.asyncio +async def test_postgres_reconciliation_uses_metadata_and_atomically_claims_expired_run(postgres_repositories): + task_repo, task_run_repo, run_repo = postgres_repositories + now = datetime.now(UTC) + for suffix in ("live", "expired"): + await _create_cron_task(task_repo, f"task-{suffix}", next_run_at=None) + await task_run_repo.create( + run_record_id=f"task-run-{suffix}", + task_id=f"task-{suffix}", + thread_id=f"thread-{suffix}", + scheduled_for=now, + trigger="scheduled", + status="running" if suffix == "expired" else "queued", + ) + + await task_run_repo.update_status("task-run-expired", status="running", run_id="run-expired") + await run_repo.put( + "run-live", + thread_id="thread-live", + user_id="user-1", + status="running", + metadata={ + "scheduled_task_id": "task-live", + "scheduled_task_run_id": "task-run-live", + }, + owner_worker_id="pod-a", + lease_expires_at=(now + timedelta(seconds=60)).isoformat(), + ) + await run_repo.put( + "run-expired", + thread_id="thread-expired", + user_id="user-1", + status="running", + metadata={ + "scheduled_task_id": "task-expired", + "scheduled_task_run_id": "task-run-expired", + }, + owner_worker_id="pod-dead", + lease_expires_at=(now - timedelta(seconds=60)).isoformat(), + ) + + assert await task_run_repo.reconcile_active_runs(error="lease expired", now=now) == 1 + assert (await task_run_repo.list_by_task("task-live"))[0]["status"] == "queued" + assert (await task_run_repo.list_by_task("task-expired"))[0]["status"] == "interrupted" + recovered = await run_repo.get("run-expired", user_id=None) + assert recovered is not None + assert recovered["status"] == "error" + assert recovered["stop_reason"] == "scheduled_task_orphan_recovered" diff --git a/backend/tests/test_scheduled_task_repository.py b/backend/tests/test_scheduled_task_repository.py index a9101f440..92f11bd36 100644 --- a/backend/tests/test_scheduled_task_repository.py +++ b/backend/tests/test_scheduled_task_repository.py @@ -1,10 +1,11 @@ -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta import pytest from deerflow.config.database_config import DatabaseConfig from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config -from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository +from deerflow.persistence.run import RunRepository +from deerflow.persistence.scheduled_task_runs import ActiveScheduledRunConflict, ScheduledTaskRunRepository from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository @@ -110,6 +111,167 @@ async def test_mark_stale_active_runs_fails_orphaned_runs(tmp_path): await close_engine() +@pytest.mark.asyncio +async def test_lease_aware_recovery_preserves_live_peer_and_reclaims_expired_peer(tmp_path): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + task_repo = ScheduledTaskRepository(sf) + task_run_repo = ScheduledTaskRunRepository(sf) + durable_run_repo = RunRepository(sf) + now = datetime.now(UTC) + + for suffix in ("live", "expired"): + await task_repo.create( + task_id=f"task-{suffix}", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title=suffix, + prompt="p", + schedule_type="cron", + schedule_spec={"cron": "* * * * *"}, + timezone="UTC", + next_run_at=None, + ) + await task_run_repo.create( + run_record_id=f"task-run-{suffix}", + task_id=f"task-{suffix}", + thread_id=f"thread-{suffix}", + scheduled_for=now, + trigger="scheduled", + status="running", + ) + await task_run_repo.update_status(f"task-run-{suffix}", status="running", run_id=f"run-{suffix}") + + await durable_run_repo.put( + "run-live", + thread_id="thread-live", + user_id="user-1", + status="running", + owner_worker_id="worker-a", + lease_expires_at=(now + timedelta(seconds=60)).isoformat(), + ) + await durable_run_repo.put( + "run-expired", + thread_id="thread-expired", + user_id="user-1", + status="running", + owner_worker_id="worker-dead", + lease_expires_at=(now - timedelta(seconds=60)).isoformat(), + ) + + reconciled = await task_run_repo.reconcile_active_runs(error="restart", now=now) + assert reconciled == 1 + assert (await task_run_repo.list_by_task("task-live"))[0]["status"] == "running" + assert (await task_run_repo.list_by_task("task-expired"))[0]["status"] == "interrupted" + recovered = await durable_run_repo.get("run-expired", user_id=None) + assert recovered is not None + assert recovered["status"] == "error" + assert recovered["stop_reason"] == "scheduled_task_orphan_recovered" + with pytest.raises(ActiveScheduledRunConflict): + await task_run_repo.create( + run_record_id="task-run-live-duplicate", + task_id="task-live", + thread_id="thread-new", + scheduled_for=now, + trigger="scheduled", + status="queued", + ) + finally: + await close_engine() + + +@pytest.mark.asyncio +async def test_lease_aware_recovery_preserves_queued_dispatch_until_lease_expires(tmp_path): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + task_repo = ScheduledTaskRepository(sf) + task_run_repo = ScheduledTaskRunRepository(sf) + now = datetime.now(UTC) + await task_repo.create( + task_id="task-queued", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="queued", + prompt="p", + schedule_type="cron", + schedule_spec={"cron": "* * * * *"}, + timezone="UTC", + next_run_at=None, + ) + assert await task_repo.claim_dispatch_lease("task-queued", lease_owner="worker-a", now=now, lease_seconds=120) is not None + assert await task_repo.claim_dispatch_lease("task-queued", lease_owner="worker-b", now=now, lease_seconds=120) is None + await task_run_repo.create( + run_record_id="task-run-queued-live", + task_id="task-queued", + thread_id="thread-queued", + scheduled_for=now, + trigger="manual", + status="queued", + ) + + assert await task_run_repo.reconcile_active_runs(error="restart", now=now) == 0 + assert (await task_run_repo.list_by_task("task-queued"))[0]["status"] == "queued" + assert await task_run_repo.reconcile_active_runs(error="restart", now=now + timedelta(seconds=121)) == 1 + assert (await task_run_repo.list_by_task("task-queued"))[0]["status"] == "interrupted" + finally: + await close_engine() + + +@pytest.mark.asyncio +async def test_lease_aware_recovery_uses_parent_last_run_when_row_link_is_missing(tmp_path): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + task_repo = ScheduledTaskRepository(sf) + task_run_repo = ScheduledTaskRunRepository(sf) + durable_run_repo = RunRepository(sf) + now = datetime.now(UTC) + await task_repo.create( + task_id="task-missing-link", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="missing link", + prompt="p", + schedule_type="cron", + schedule_spec={"cron": "* * * * *"}, + timezone="UTC", + next_run_at=None, + ) + await task_repo.update("task-missing-link", user_id="user-1", updates={"last_run_id": "run-peer"}) + await task_run_repo.create( + run_record_id="task-run-missing-link", + task_id="task-missing-link", + thread_id="thread-peer", + scheduled_for=now, + trigger="scheduled", + status="running", + ) + await durable_run_repo.put( + "run-peer", + thread_id="thread-peer", + user_id="user-1", + status="running", + owner_worker_id="worker-a", + lease_expires_at=(now + timedelta(seconds=60)).isoformat(), + ) + + assert await task_run_repo.reconcile_active_runs(error="restart", now=now) == 0 + assert (await task_run_repo.list_by_task("task-missing-link"))[0]["status"] == "running" + finally: + await close_engine() + + @pytest.mark.asyncio async def test_update_status_protect_terminal_keeps_completion_result(tmp_path): """The launch-path "running" write must not clobber a terminal status @@ -220,6 +382,95 @@ async def test_cancel_stuck_once_tasks_reconciles_orphaned_running(tmp_path): await close_engine() +@pytest.mark.asyncio +async def test_lease_aware_once_recovery_keeps_live_peer_and_cancels_dead_run(tmp_path): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + task_repo = ScheduledTaskRepository(sf) + durable_run_repo = RunRepository(sf) + now = datetime.now(UTC) + for suffix in ("live", "dead"): + await task_repo.create( + task_id=f"task-once-{suffix}", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title=suffix, + prompt="p", + schedule_type="once", + schedule_spec={"run_at": (now + timedelta(minutes=5)).isoformat()}, + timezone="UTC", + next_run_at=None, + ) + await task_repo.update( + f"task-once-{suffix}", + user_id="user-1", + updates={"status": "running", "last_run_id": f"run-once-{suffix}"}, + ) + + await durable_run_repo.put( + "run-once-live", + thread_id="thread-live", + user_id="user-1", + status="running", + owner_worker_id="worker-a", + lease_expires_at=(now + timedelta(seconds=60)).isoformat(), + ) + await durable_run_repo.put( + "run-once-dead", + thread_id="thread-dead", + user_id="user-1", + status="error", + owner_worker_id="worker-dead", + lease_expires_at=(now - timedelta(seconds=60)).isoformat(), + ) + + assert await task_repo.reconcile_stuck_once_tasks(error="restart", now=now) == 1 + live = await task_repo.get("task-once-live", user_id="user-1") + dead = await task_repo.get("task-once-dead", user_id="user-1") + assert live is not None and live["status"] == "running" + assert dead is not None and dead["status"] == "cancelled" + finally: + await close_engine() + + +@pytest.mark.asyncio +async def test_lease_aware_once_recovery_reclaims_expired_dispatch_lease(tmp_path): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + repo = ScheduledTaskRepository(sf) + now = datetime.now(UTC) + await repo.create( + task_id="task-once-expired-lease", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="expired lease", + prompt="p", + schedule_type="once", + schedule_spec={"run_at": now.isoformat()}, + timezone="UTC", + next_run_at=now, + ) + await repo.update( + "task-once-expired-lease", + user_id="user-1", + updates={"status": "running", "lease_expires_at": now - timedelta(seconds=60)}, + ) + + assert await repo.reconcile_stuck_once_tasks(error="restart", now=now) == 1 + task = await repo.get("task-once-expired-lease", user_id="user-1") + assert task is not None and task["status"] == "cancelled" + finally: + await close_engine() + + @pytest.mark.asyncio async def test_update_after_launch_protect_terminal_keeps_hook_result(tmp_path): """The launch-path bookkeeping write must not clobber a terminal task @@ -322,3 +573,266 @@ async def test_list_by_user_and_thread_filters_in_sql(tmp_path): assert await repo.list_by_user_and_thread("user-2", "thread-1") == [] await close_engine() + + +@pytest.mark.asyncio +async def test_reconcile_recovers_live_run_link_from_metadata(tmp_path): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + task_repo = ScheduledTaskRepository(sf) + task_run_repo = ScheduledTaskRunRepository(sf) + run_repo = RunRepository(sf) + now = datetime.now(UTC) + + await task_repo.create( + task_id="task-metadata-link", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="metadata fallback", + prompt="p", + schedule_type="cron", + schedule_spec={"cron": "* * * * *"}, + timezone="UTC", + next_run_at=None, + ) + await task_run_repo.create( + run_record_id="task-run-metadata-link", + task_id="task-metadata-link", + thread_id="thread-metadata-link", + scheduled_for=now, + trigger="scheduled", + status="queued", + ) + await run_repo.put( + "run-metadata-link", + thread_id="thread-metadata-link", + user_id="user-1", + status="running", + metadata={ + "scheduled_task_id": "task-metadata-link", + "scheduled_task_run_id": "task-run-metadata-link", + }, + owner_worker_id="worker-a", + lease_expires_at=(now + timedelta(seconds=60)).isoformat(), + ) + + assert await task_run_repo.reconcile_active_runs(error="lease expired", now=now) == 0 + row = (await task_run_repo.list_by_task("task-metadata-link"))[0] + assert row["status"] == "queued" + assert row["run_id"] is None + finally: + await close_engine() + + +@pytest.mark.asyncio +async def test_reconcile_ignores_stale_parent_last_run_before_metadata_fallback(tmp_path): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + task_repo = ScheduledTaskRepository(sf) + task_run_repo = ScheduledTaskRunRepository(sf) + run_repo = RunRepository(sf) + now = datetime.now(UTC) + + await task_repo.create( + task_id="task-stale-parent-link", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="stale parent link", + prompt="p", + schedule_type="cron", + schedule_spec={"cron": "* * * * *"}, + timezone="UTC", + next_run_at=None, + ) + await task_run_repo.create( + run_record_id="task-run-current", + task_id="task-stale-parent-link", + thread_id="thread-current", + scheduled_for=now, + trigger="scheduled", + status="queued", + ) + await task_repo.update("task-stale-parent-link", user_id="user-1", updates={"last_run_id": "run-previous"}) + await run_repo.put( + "run-previous", + thread_id="thread-previous", + user_id="user-1", + status="success", + metadata={"scheduled_task_id": "task-stale-parent-link", "scheduled_task_run_id": "task-run-previous"}, + ) + await run_repo.put( + "run-current", + thread_id="thread-current", + user_id="user-1", + status="running", + metadata={ + "scheduled_task_id": "task-stale-parent-link", + "scheduled_task_run_id": "task-run-current", + }, + owner_worker_id="worker-a", + lease_expires_at=(now + timedelta(seconds=60)).isoformat(), + ) + + assert await task_run_repo.reconcile_active_runs(error="lease expired", now=now) == 0 + assert (await task_run_repo.list_by_task("task-stale-parent-link"))[0]["status"] == "queued" + finally: + await close_engine() + + +@pytest.mark.asyncio +async def test_reconcile_preserves_row_when_heartbeat_wins_takeover(tmp_path): + class RenewedRunRepository: + async def claim_for_takeover(self, *_args, **_kwargs): + return False + + async def get(self, *_args, **_kwargs): + return {"status": "running"} + + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + task_repo = ScheduledTaskRepository(sf) + task_run_repo = ScheduledTaskRunRepository(sf, run_repository=RenewedRunRepository()) + run_repo = RunRepository(sf) + now = datetime.now(UTC) + await task_repo.create( + task_id="task-heartbeat-race", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="heartbeat race", + prompt="p", + schedule_type="cron", + schedule_spec={"cron": "* * * * *"}, + timezone="UTC", + next_run_at=None, + ) + await task_run_repo.create( + run_record_id="task-run-heartbeat-race", + task_id="task-heartbeat-race", + thread_id="thread-heartbeat-race", + scheduled_for=now, + trigger="scheduled", + status="running", + ) + await task_run_repo.update_status( + "task-run-heartbeat-race", + status="running", + run_id="run-heartbeat-race", + ) + await run_repo.put( + "run-heartbeat-race", + thread_id="thread-heartbeat-race", + user_id="user-1", + status="running", + owner_worker_id="worker-a", + lease_expires_at=(now - timedelta(seconds=60)).isoformat(), + ) + + assert await task_run_repo.reconcile_active_runs(error="lease expired", now=now) == 0 + row = (await task_run_repo.list_by_task("task-heartbeat-race"))[0] + assert row["status"] == "running" + finally: + await close_engine() + + +@pytest.mark.asyncio +async def test_update_after_launch_rejects_stale_lease_owner(tmp_path, caplog): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + repo = ScheduledTaskRepository(sf) + now = datetime.now(UTC) + await repo.create( + task_id="task-fenced", + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title="fenced", + prompt="p", + schedule_type="cron", + schedule_spec={"cron": "* * * * *"}, + timezone="UTC", + next_run_at=now, + ) + assert await repo.claim_dispatch_lease("task-fenced", lease_owner="worker-a", now=now, lease_seconds=60) is not None + await repo.update( + "task-fenced", + user_id="user-1", + updates={ + "lease_owner": "worker-b", + "lease_expires_at": now + timedelta(seconds=120), + }, + ) + + with caplog.at_level("WARNING", logger="deerflow.persistence.scheduled_tasks.sql"): + updated = await repo.update_after_launch( + "task-fenced", + status="enabled", + next_run_at=now + timedelta(minutes=1), + last_run_at=now, + last_run_id="run-a", + last_thread_id="thread-a", + last_error=None, + increment_run_count=True, + expected_lease_owner="worker-a", + ) + assert updated is False + assert "task-fenced" in caplog.text + assert "expected lease owner worker-a, current owner worker-b" in caplog.text + task = await repo.get("task-fenced", user_id="user-1") + assert task is not None + assert task["lease_owner"] == "worker-b" + assert task["last_run_id"] is None + assert task["run_count"] == 0 + finally: + await close_engine() + + +@pytest.mark.asyncio +async def test_global_budget_counts_dispatch_lease_reservations(tmp_path): + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + sf = get_session_factory() + assert sf is not None + repo = ScheduledTaskRepository(sf) + now = datetime.now(UTC) + for task_id in ("task-reserved", "task-due-a", "task-due-b"): + await repo.create( + task_id=task_id, + user_id="user-1", + thread_id=None, + context_mode="fresh_thread_per_run", + assistant_id="lead_agent", + title=task_id, + prompt="p", + schedule_type="cron", + schedule_spec={"cron": "* * * * *"}, + timezone="UTC", + next_run_at=None if task_id == "task-reserved" else now, + ) + assert await repo.claim_dispatch_lease("task-reserved", lease_owner="worker-a", now=now, lease_seconds=60) is not None + + claimed = await repo.claim_due_tasks( + now=now, + lease_owner="worker-b", + lease_seconds=60, + limit=2, + global_max_concurrent_runs=2, + ) + assert len(claimed) == 1 + assert claimed[0]["id"] in {"task-due-a", "task-due-b"} + finally: + await close_engine() diff --git a/backend/tests/test_scheduled_task_service.py b/backend/tests/test_scheduled_task_service.py index cf185dd36..cafc33bf8 100644 --- a/backend/tests/test_scheduled_task_service.py +++ b/backend/tests/test_scheduled_task_service.py @@ -1,3 +1,4 @@ +import asyncio from datetime import UTC, datetime, timedelta import pytest @@ -14,11 +15,19 @@ class DummyTaskRepo: self.claimed = False self.updated = None self.cancelled_stuck_once = None + self.reconciled_stuck_once = None async def cancel_stuck_once_tasks(self, *, error): self.cancelled_stuck_once = error return 0 + async def reconcile_stuck_once_tasks(self, **kwargs): + self.reconciled_stuck_once = kwargs + return 0 + + async def claim_dispatch_lease(self, task_id, **_kwargs): + return next((dict(row) for row in self.rows if row["id"] == task_id), None) + async def claim_due_tasks(self, **_kwargs): if self.claimed: return [] @@ -47,6 +56,8 @@ class DummyRunRepo: self.active = active self.active_count = active_count self.stale_marked = None + self.reconciled = None + self.reconcile_count = 0 async def count_active_runs(self): return self.active_count @@ -65,6 +76,11 @@ class DummyRunRepo: self.stale_marked = error return 0 + async def reconcile_active_runs(self, **kwargs): + self.reconcile_count += 1 + self.reconciled = kwargs + return 0 + @pytest.mark.asyncio async def test_service_claims_and_dispatches_due_task(): @@ -576,6 +592,33 @@ async def test_startup_sweep_reconciles_stale_runs_and_stuck_once_tasks(): assert task_repo.cancelled_stuck_once == run_repo.stale_marked +@pytest.mark.asyncio +async def test_multi_instance_start_uses_lease_aware_reconciliation(): + task_repo = DummyTaskRepo([]) + run_repo = DummyRunRepo() + service = ScheduledTaskService( + task_repo=task_repo, + task_run_repo=run_repo, + launch_run=lambda **_kwargs: None, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_runs=3, + multi_instance=True, + run_lease_grace_seconds=17, + ) + + await service.start() + await asyncio.sleep(0) + await service.stop() + + assert run_repo.reconcile_count == 1 + assert run_repo.reconciled is not None + assert run_repo.reconciled["lease_grace_seconds"] == 17 + assert task_repo.reconciled_stuck_once is not None + assert task_repo.reconciled_stuck_once["lease_grace_seconds"] == 17 + assert task_repo.cancelled_stuck_once is None + + @pytest.mark.asyncio async def test_manual_trigger_with_active_run_returns_conflict_without_launching(): launched = [] @@ -668,12 +711,12 @@ class _StatefulRunRepo: _ACTIVE = {"queued", "running"} - def __init__(self, *, fail_first_update: bool = False) -> None: + def __init__(self, *, fail_first_update: bool = False, fail_updates: int = 0) -> None: self.created: list[dict] = [] self.updates: list[tuple[str, dict]] = [] self.rows: dict[str, dict] = {} - self._fail_first_update = fail_first_update - self._first_update_raised = False + self._fail_updates = max(fail_updates, 1 if fail_first_update else 0) + self._updates_raised = 0 async def count_active_runs(self) -> int: return sum(1 for row in self.rows.values() if row["status"] in self._ACTIVE) @@ -689,10 +732,11 @@ class _StatefulRunRepo: async def update_status(self, run_record_id: str, **kwargs) -> None: self.updates.append((run_record_id, kwargs)) - if self._fail_first_update and not self._first_update_raised: - # The launch-path queued->running write fails once, AFTER - # _launch_run has already returned a live run_id. - self._first_update_raised = True + if self._updates_raised < self._fail_updates: + # The launch-path queued->running write fails AFTER _launch_run has + # already returned a live run_id. Some tests fail both attempts to + # pin the last-resort active-slot behavior. + self._updates_raised += 1 raise RuntimeError("simulated transient DB error on queued->running write") row = self.rows.get(run_record_id) if row is None: @@ -784,6 +828,64 @@ async def test_post_launch_bookkeeping_failure_does_not_release_active_slot(): assert task_repo.updated[1]["last_error"] is None +@pytest.mark.asyncio +async def test_both_post_launch_association_writes_can_fail_without_releasing_slot(): + launched: list[dict] = [] + + async def fake_launch(**kwargs): + launched.append(kwargs) + return {"run_id": "run-live", "thread_id": kwargs["thread_id"]} + + class FailingTaskRepo(DummyTaskRepo): + def __init__(self, rows): + super().__init__(rows) + self.failures_remaining = 1 + + async def update_after_launch(self, *args, **kwargs): + if self.failures_remaining: + self.failures_remaining -= 1 + raise RuntimeError("simulated parent bookkeeping failure") + await super().update_after_launch(*args, **kwargs) + + task = { + "id": "task-double-failure", + "user_id": "user-1", + "thread_id": None, + "context_mode": "fresh_thread_per_run", + "assistant_id": "lead_agent", + "prompt": "do the thing", + "schedule_type": "cron", + "schedule_spec": {"cron": "*/5 * * * *"}, + "timezone": "UTC", + "status": "enabled", + "overlap_policy": "skip", + } + task_repo = FailingTaskRepo([task]) + run_repo = _StatefulRunRepo(fail_updates=2) + service = ScheduledTaskService( + task_repo=task_repo, + task_run_repo=run_repo, + launch_run=fake_launch, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_runs=3, + ) + now = datetime.now(UTC) + + first = await service.dispatch_task(dict(task), now=now, trigger="scheduled") + assert first["outcome"] == "launched" + first_row_id = run_repo.created[0]["run_record_id"] + assert run_repo.rows[first_row_id] == { + "task_id": "task-double-failure", + "status": "queued", + "run_id": None, + } + + second = await service.dispatch_task(dict(task), now=now, trigger="scheduled") + assert len(launched) == 1 + assert second["outcome"] == "skipped" + + @pytest.mark.asyncio async def test_pre_launch_failure_still_releases_active_slot(): """Complement to the #4452 fix: when ``_launch_run`` itself fails (no run diff --git a/config.example.yaml b/config.example.yaml index 5ce5cf143..9809fea5b 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -15,7 +15,7 @@ # ============================================================================ # Bump this number when the config schema changes. # Run `make config-upgrade` to merge new fields into your local config.yaml. -config_version: 33 +config_version: 34 # ============================================================================ # Logging @@ -2054,19 +2054,19 @@ agent_storage: # Background scheduler for one-time and recurring (cron) agent runs. # All fields are restart-required (captured at Gateway lifespan startup). # -# Multi-worker note: the scheduler runs once per uvicorn worker. SQLite silently -# ignores row-level locks, so multiple workers can double-fire the same task. -# For multi-worker deployments (GATEWAY_WORKERS > 1), use the Postgres database -# backend, where FOR UPDATE SKIP LOCKED serializes claims correctly. +# The scheduler is single-instance by default. Set multi_instance: true only +# when every Gateway instance shares Postgres, run ownership heartbeats, and +# database-backed run events; startup recovery then preserves live peer runs. # # scheduler: # enabled: false # Master switch for the background poller # poll_interval_seconds: 5 # How often to scan for due tasks # lease_seconds: 120 # Claim lease; a crashed process's task becomes reclaimable after this -# max_concurrent_runs: 3 # Global cap on active scheduled runs; each poll claims only into the remaining budget +# max_concurrent_runs: 3 # Global cap on active scheduled runs/reservations across multi-instance Pods # min_once_delay_seconds: 60 # Minimum future offset for one-time tasks at creation time scheduler: enabled: false + multi_instance: false # Opt into lease-aware recovery across Gateway instances poll_interval_seconds: 5 lease_seconds: 120 max_concurrent_runs: 3 diff --git a/deploy/helm/deer-flow/README.md b/deploy/helm/deer-flow/README.md index b37c75746..bd93bd3ea 100644 --- a/deploy/helm/deer-flow/README.md +++ b/deploy/helm/deer-flow/README.md @@ -124,7 +124,7 @@ they resolve from the `secrets` map): ```yaml config: | - config_version: 33 + config_version: 34 models: - name: gpt-4 use: langchain_openai:ChatOpenAI @@ -240,6 +240,15 @@ kubectl -n deer-flow exec deploy/deer-flow-provisioner -- curl -s localhost:8002 double-submit can create two runs on one thread (checkpoint corruption), a cancel can land on a non-owner pod (409), and a crashed pod's runs stay `pending`/`running` forever. Stay on 1 replica until that work lands. +- **Scheduled task recovery.** If a deployment explicitly enables + `scheduler.multi_instance: true`, it must use shared Postgres, + `run_ownership.heartbeat_enabled: true`, and `run_events.backend: db`. + Scheduler startup then preserves live scheduled runs owned by another Pod, + atomically takes over only expired leases, and fences stale post-launch + bookkeeping. `max_concurrent_runs` is a shared global cap across Pods, + including pre-launch dispatch reservations. Restart all Gateway Pods after + changing these startup-only settings. This does not remove the broader + Gateway replica limitations described above. - **Redis stream bridge.** A bundled single-instance redis StatefulSet (`redis.enabled: true`, `redis:7-alpine`) runs in the namespace and the gateway connects via the in-cluster Service. Per-run SSE events are stored in diff --git a/deploy/helm/deer-flow/values.yaml b/deploy/helm/deer-flow/values.yaml index 6dafba56e..6bf3d47e5 100644 --- a/deploy/helm/deer-flow/values.yaml +++ b/deploy/helm/deer-flow/values.yaml @@ -27,7 +27,9 @@ image: gateway: replicas: 1 # Safe default. Postgres + the Redis stream bridge are # wired, but multi-replica needs issue #3948's run-control - # work (cancel/dedup/reconcile) — see README "Gateway replicas". + # work (cancel/dedup/reconcile). Scheduler multi_instance + # recovery is safe only when its config prerequisites are met; + # max_concurrent_runs is shared across Pods in that mode. # Pod termination grace period (seconds). MUST exceed the Gateway's graceful # shutdown work: channel stop (~5s) + memory queue drain # (memory.shutdown_flush_timeout_seconds, default 30s) + buffer. K8s defaults @@ -240,7 +242,7 @@ ingress: # -- DeerFlow config.yaml content. Secrets MUST stay as $VAR references — never # inline literal secret values here. The default enables provisioner sandbox. config: | - config_version: 33 + config_version: 34 log_level: info models: []