mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-09 21:49:37 +00:00
fix(scheduler): enqueue busy scheduled task runs (#4918)
* fix(frontend): clarify reuse-thread scheduling behavior * fix(scheduler): enqueue overlapping scheduled runs * fix(scheduler): preserve queue lease fencing * fix(scheduler): close queue concurrency races * fix(scheduler): harden queue timeout bookkeeping * fix(scheduler): preserve manual failure schedule --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
336cd3acc4
commit
645ca08f16
@ -104,6 +104,7 @@ Skill quality review note:
|
||||
Scheduled-task note:
|
||||
- The scheduled-task MVP adds a workspace page at `/workspace/scheduled-tasks` plus a background scheduler service gated by `config.yaml -> scheduler.enabled`.
|
||||
- Scheduled background runs are intentionally non-interactive: they execute through the normal run lifecycle, but the lead-agent toolset excludes `ask_clarification` when `context.non_interactive=true`. The key is honored only for internally-authenticated callers (the scheduler launch path); client-supplied `context.non_interactive` is dropped.
|
||||
- Busy scheduled occurrences are persisted as `queued`; `launching` is a short lease-fenced claim, `running` remains the normal Gateway run lifecycle, and `scheduler.queue_timeout_seconds` bounds the durable wait. Do not reintroduce skip-on-overlap or count waiting rows against `max_concurrent_runs`.
|
||||
- Scheduled launches use `scheduler.recursion_limit` (default 1000, matching the web UI's `recursion_limit: 1000`, clamped by `max_recursion_limit`). The value is read at dispatch, so a YAML edit applies to the next scheduled run without a Gateway restart.
|
||||
|
||||
## Commands: Root vs. Module
|
||||
|
||||
@ -1347,10 +1347,11 @@ DeerFlow now includes a first-class scheduled-task MVP in the workspace.
|
||||
Current MVP capabilities:
|
||||
|
||||
- Manage tasks at `/workspace/scheduled-tasks`
|
||||
- Choose whether each scheduled task reuses a thread or creates a fresh thread per run
|
||||
- Choose whether each scheduled task reuses a thread and its conversation history or creates a fresh thread per run
|
||||
- Support `once` and `cron` schedules
|
||||
- Run background scheduled executions as non-interactive DeerFlow runs (`ask_clarification` is not exposed there)
|
||||
- Use `skip` overlap behavior for due cron executions that collide with an active run on the same reused thread
|
||||
- Persist a due execution as `queued` when its reused thread or the global execution budget is busy, then launch it when capacity is available; queued occurrences survive Gateway restarts and fail after `scheduler.queue_timeout_seconds`
|
||||
- Freeze a task's definition while an occurrence is `queued`, `launching`, or `running`, so a durable occurrence cannot silently pick up a different prompt, thread, or schedule; transitioning a task to paused or deleting it cancels an existing waiting occurrence, while `launching`/`running` work must finish before those mutations are retried and an explicit manual trigger may still wait and run without resuming a paused schedule
|
||||
- Pause, resume, trigger, inspect history, and delete tasks
|
||||
- Execute scheduled work through the normal DeerFlow run lifecycle
|
||||
|
||||
@ -1365,12 +1366,12 @@ Enable background polling with `config.yaml -> scheduler.enabled`. Manual trigge
|
||||
|
||||
Scheduled runs use `scheduler.recursion_limit` in `config.yaml` (default `1000`, matching the web UI's interactive budget). Values above `max_recursion_limit` are clamped. This field is read at dispatch, so the next scheduled run picks it up without a Gateway restart.
|
||||
|
||||
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.
|
||||
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 return expired launch claims to the queue, take over only expired run leases, and fence stale launch writes. `max_concurrent_runs` is a shared global cap across Pods for `launching`/`running` occurrences; waiting `queued` rows do not consume it. 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.
|
||||
- In multi-instance mode, `scheduler.max_concurrent_runs` is a cluster-wide execution cap, not a per-Pod cap. It includes `launching` and `running` scheduled occurrences, so capacity does not multiply with the number of replicas; durable waiting rows remain outside the cap.
|
||||
- `scheduler.multi_instance` and the related scheduler, ownership, and run-event settings are startup-only. Apply changes with a coordinated restart of all Gateway Pods; changing the ConfigMap alone does not activate multi-instance recovery.
|
||||
|
||||
## Terminal Workbench (TUI)
|
||||
|
||||
@ -725,10 +725,11 @@ DeerFlow 现在在 workspace 里内置了一个一等的定时任务(scheduled
|
||||
当前 MVP 能力:
|
||||
|
||||
- 在 `/workspace/scheduled-tasks` 管理任务
|
||||
- 每个定时任务可以选择复用同一个 thread,也可以选择每次运行新建一个 thread
|
||||
- 每个定时任务可以选择复用同一个 thread 及其历史对话,也可以选择每次运行新建一个 thread
|
||||
- 支持 `once` 和 `cron` 两种调度方式
|
||||
- 后台定时执行以非交互式 DeerFlow run 运行(那里不会暴露 `ask_clarification`)
|
||||
- 当到期的 cron 执行与同一复用 thread 上的活跃 run 冲突时,采用 `skip` 的重叠处理策略
|
||||
- 当所复用的 thread 或全局执行配额正忙时,到期执行会持久化为 `queued`,并在可用后启动;队列项在 Gateway 重启后保留,超过 `scheduler.queue_timeout_seconds` 后标记为失败
|
||||
- 当某次执行处于 `queued`、`launching` 或 `running` 时冻结任务定义,避免持久化的执行意外换用新的 prompt、thread 或调度;将任务切换为暂停或删除任务会取消已在等待的执行,而 `launching`/`running` 执行结束后才能重试这些变更;显式手动触发在调度已暂停时仍可等待并执行,且不会自动恢复调度
|
||||
- 支持暂停、恢复、手动触发、查看历史和删除任务
|
||||
- 定时任务通过正常的 DeerFlow run 生命周期执行
|
||||
|
||||
|
||||
@ -18,11 +18,11 @@ 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|<ns>`, 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. Scheduled launches pass `scheduler.recursion_limit` (default 1000, matching the web UI's `recursion_limit: 1000`, clamped by `max_recursion_limit`) via `launch_scheduled_thread_run`; the value is read from `get_app_config()` at dispatch.
|
||||
- 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).
|
||||
- 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 launch claims return to the durable queue, expired run leases are atomically taken over, stale launch writes are fenced by lease ownership, and the Postgres advisory-locked budget makes `max_concurrent_runs` a shared global cap for `launching`/`running` rows.
|
||||
- Long-running MCP work uses a separate durable task runtime rather than keeping remote task IDs or status polling inside the Agent loop. Explicit `task_toolsets` bind raw submit/status/cancel names; only submit remains Agent-visible, and its wrapper persists the remote handle before returning a local ID. `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 or after a cancel request must be discarded even when the owner token still matches. The first cancel request fences an in-flight poll lease, while repeats preserve an active cancellation lease so they cannot issue concurrent remote cancels; cancellation backoff starts when the remote attempt finishes, so a slow timeout cannot consume the retry delay. Cancellation, polling, and notification batches isolate per-task exceptions; an unexpected cancellation/poll failure leaves that record's lease to expire, while notification failures release only the affected lease for retry. Input-required and terminal event snapshots are delivered by idempotent Agent runs and marked delivered only after run success; the trusted notification instruction stays outside the input boundary while the serialized remote event is framed as untrusted data. A busy-thread conflict is normalized back to the service boundary so the queued snapshot coalesces to the latest task event. A missing dispatched run becomes a failed delivery attempt, while transient run-store hydration errors stay distinguishable and retry the same lookup. The database is the source of truth; `ThreadState` receives only a bounded current-thread projection, and display names are neutralized at that model-state boundary. The installed process-local submitter is the source of truth for management-tool exposure; hot `mcp_tasks` edits take effect only after restart, and active skills must explicitly declare the list/cancel business tools.
|
||||
- MCP notification failures use a consecutive counter separate from the idempotency-key `dispatch_attempt`, capped exponential backoff, latest-event rebuilding before a run launches, and a five-attempt budget before `dead_letter`. A permanently missing/mismatched target thread is dead-lettered immediately instead of being recreated or reclaimed. HTTP and Agent cancellation requests return after the durable cancel fence; the background loop alone owns the potentially slow remote call and retry schedule. The HTTP cancel endpoint rejects requests with 503 when the loop is not running (`mcp_tasks_available` false, e.g. `mcp_tasks.enabled=false` with SQL persistence), so a cancellation is never acknowledged without a worker to perform it. The bounded notification error/count/status join poll and cancellation diagnostics in the task detail API and expanded card.
|
||||
- Scheduled-task dispatch enforces at most one non-terminal occurrence per task through `uq_scheduled_task_run_active` (`task_id WHERE status IN ('queued','launching','running')`). `queued` is durable and survives restart; `launching` carries a short owner/expiry lease and is the only state that may call the normal Gateway launch path; `running` is associated with the durable run. Each occurrence also supplies a stable run-admission idempotency key, so a recovered launch retry reuses the same durable run. A reused-thread `ConflictError` moves `launching` back to `queued`, while non-conflict launch errors become terminal `failed`. Waiting rows do not consume `max_concurrent_runs`; the atomic queue claim enforces the budget. Repeated triggers coalesce on the one active row, and same-thread FIFO treats older `queued`, `launching`, and `running` rows as blockers. The task definition stays immutable for all three active states because queue admission, PATCH/resume, pause, and delete serialize on the parent task row before touching the occurrence row. Pause/delete atomically interrupt existing `queued` rows and reject `launching`/`running` rows; PATCH/resume reject every active state, and mutation errors advertise pause cancellation only for `queued` work. A manual trigger may queue and run while the parent schedule remains paused. Recovery and multi-instance reconciliation lock task/run pairs in deterministic task-id/run-id order and must reconstruct `run_id`, `started_at`, and the live error state before releasing the short launch claim. Launch/failure/timeout bookkeeping changes the occurrence and its parent task in one parent-first transaction so a peer cannot claim the released task between those writes. Queue timeout marks the occurrence failed and advances a scheduled occurrence so it cannot immediately requeue forever; repository write boundaries coerce serialized task timestamps before binding SQL `DateTime` fields.
|
||||
- `extensions_config.json` is written at runtime by the Gateway (`PUT`/`PATCH /api/mcp/config`, the MCP enable switch, skill updates), so the production compose mounts it read-write while `config.yaml` stays `:ro`; Helm copies its ConfigMap seed into a writable home-volume directory before Gateway starts. Every read-modify-write holds both `extensions_config_write_lock` and the sidecar advisory `extensions_config_file_lock`, because the process-local lock alone loses updates across workers. Docker mounts the compose file as its own mount point, and Linux refuses `rename()` over a mount point with `EBUSY` even when the mount is writable — so `atomic_write_extensions_config` keeps the temp-file-plus-rename path and falls back to an in-place overwrite only on `EBUSY`. That fallback is deliberately non-atomic (a crash mid-write truncates the file); it exists because the alternative is a write that can never succeed, and only its first occurrence per target is logged at warning level. Any other `errno` still propagates. Pinned by `tests/test_compose_extensions_config_writable.py`, `tests/test_extensions_config_atomic_write.py`, and `tests/test_helm_extensions_config_writable.py`.
|
||||
- Scheduled-task dispatch enforces "at most one active run per task when `overlap_policy=skip`" at the DB layer via the partial unique index `uq_scheduled_task_run_active` (`scheduled_task_runs.task_id WHERE status IN ('queued','running')`). `ScheduledTaskService.dispatch_task`'s `has_active_runs` check is a non-atomic fast path (its own session, separated from the `create()` insert by `await` points), so two concurrent dispatches — a manual `POST /scheduled-tasks/{id}/trigger` racing the poller, a double-click, or a client retry — can both pass it; the index is the atomic arbiter, and the losing `create` surfaces as `ActiveScheduledRunConflict` (translated from `IntegrityError` in the repository) and collapses to the same outcome as the fast path (manual → 409 conflict, scheduled → a `"skipped"` tombstone). The scheduled-skip tombstone is created directly as terminal `"skipped"` (not a transient `"queued"`) so it never occupies the active slot the pre-existing run still holds. Sibling of the `runs` table's `uq_runs_thread_active` (PR #4003), which keys on `thread_id` and so does not cover the default `fresh_thread_per_run` context where every dispatch gets a new thread. Index is status-only, not `overlap_policy`-conditional (the policy is fixed to `"skip"` in the MVP).
|
||||
|
||||
**Project Structure**:
|
||||
```
|
||||
|
||||
@ -325,6 +325,7 @@ 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,
|
||||
queue_timeout_seconds=startup_config.scheduler.queue_timeout_seconds,
|
||||
multi_instance=startup_config.scheduler.multi_instance,
|
||||
run_lease_grace_seconds=startup_config.run_ownership.grace_seconds,
|
||||
)
|
||||
|
||||
@ -16,6 +16,7 @@ from app.gateway.deps import (
|
||||
get_scheduled_task_service,
|
||||
get_thread_store,
|
||||
)
|
||||
from deerflow.persistence.scheduled_tasks import ActiveScheduledTaskMutationConflict
|
||||
from deerflow.scheduler.schedules import (
|
||||
next_run_at as compute_next_run_at,
|
||||
)
|
||||
@ -28,12 +29,25 @@ from deerflow.utils.thread_id import ThreadId
|
||||
router = APIRouter(prefix="/api", tags=["scheduled-tasks"])
|
||||
|
||||
|
||||
def _ensure_task_mutable(task: dict[str, Any]) -> None:
|
||||
def _active_occurrence_conflict_detail(status: str) -> str:
|
||||
detail = f"Scheduled task has an active {status} occurrence; retry after it finishes"
|
||||
if status == "queued":
|
||||
detail += " or cancel the queued occurrence by pausing the task"
|
||||
return detail
|
||||
|
||||
|
||||
async def _ensure_task_mutable(task: dict[str, Any], repo) -> None:
|
||||
if task.get("status") == "running":
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Scheduled task is currently running; retry after the active execution finishes",
|
||||
)
|
||||
active_status = await repo.get_active_run_status(task["id"])
|
||||
if active_status is not None:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=_active_occurrence_conflict_detail(active_status),
|
||||
)
|
||||
|
||||
|
||||
class ScheduledTaskCreateRequest(BaseModel):
|
||||
@ -148,7 +162,7 @@ async def update_scheduled_task(task_id: str, request: Request, body: ScheduledT
|
||||
existing = await repo.get(task_id, user_id=str(user.id))
|
||||
if existing is None:
|
||||
raise HTTPException(status_code=404, detail="Scheduled task not found")
|
||||
_ensure_task_mutable(existing)
|
||||
await _ensure_task_mutable(existing, repo)
|
||||
|
||||
updates = body.model_dump(exclude_none=True)
|
||||
if "context_mode" in updates:
|
||||
@ -208,11 +222,20 @@ async def update_scheduled_task(task_id: str, request: Request, body: ScheduledT
|
||||
if next_run_at is not None and existing["status"] in {"completed", "failed", "cancelled"}:
|
||||
updates["status"] = "enabled"
|
||||
|
||||
updated = await repo.update(
|
||||
task_id,
|
||||
user_id=str(user.id),
|
||||
updates=updates,
|
||||
)
|
||||
try:
|
||||
updated = await repo.update(
|
||||
task_id,
|
||||
user_id=str(user.id),
|
||||
updates=updates,
|
||||
require_mutable=True,
|
||||
)
|
||||
except ActiveScheduledTaskMutationConflict as exc:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=_active_occurrence_conflict_detail(exc.status),
|
||||
) from exc
|
||||
if updated is None:
|
||||
raise HTTPException(status_code=404, detail="Scheduled task not found")
|
||||
return updated
|
||||
|
||||
|
||||
@ -226,11 +249,25 @@ async def pause_scheduled_task(task_id: str, request: Request):
|
||||
existing = await repo.get(task_id, user_id=str(user.id))
|
||||
if existing is None:
|
||||
raise HTTPException(status_code=404, detail="Scheduled task not found")
|
||||
_ensure_task_mutable(existing)
|
||||
updated = await repo.update(task_id, user_id=str(user.id), updates={"status": "paused"})
|
||||
if updated is None:
|
||||
if existing.get("status") == "running":
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Scheduled task is currently running; retry after the active execution finishes",
|
||||
)
|
||||
result = await repo.pause_with_queue_cancellation(
|
||||
task_id,
|
||||
user_id=str(user.id),
|
||||
error="scheduled task was paused while queued",
|
||||
now=datetime.now(UTC),
|
||||
)
|
||||
if result == "not_found":
|
||||
raise HTTPException(status_code=404, detail="Scheduled task not found")
|
||||
return updated
|
||||
if result == "executing":
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Scheduled task is already launching or running; retry after the active execution finishes",
|
||||
)
|
||||
return await repo.get(task_id, user_id=str(user.id))
|
||||
|
||||
|
||||
@router.post("/scheduled-tasks/{task_id}/resume")
|
||||
@ -243,8 +280,19 @@ async def resume_scheduled_task(task_id: str, request: Request):
|
||||
existing = await repo.get(task_id, user_id=str(user.id))
|
||||
if existing is None:
|
||||
raise HTTPException(status_code=404, detail="Scheduled task not found")
|
||||
_ensure_task_mutable(existing)
|
||||
updated = await repo.update(task_id, user_id=str(user.id), updates={"status": "enabled"})
|
||||
await _ensure_task_mutable(existing, repo)
|
||||
try:
|
||||
updated = await repo.update(
|
||||
task_id,
|
||||
user_id=str(user.id),
|
||||
updates={"status": "enabled"},
|
||||
require_mutable=True,
|
||||
)
|
||||
except ActiveScheduledTaskMutationConflict as exc:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail=_active_occurrence_conflict_detail(exc.status),
|
||||
) from exc
|
||||
if updated is None:
|
||||
raise HTTPException(status_code=404, detail="Scheduled task not found")
|
||||
return updated
|
||||
@ -262,6 +310,8 @@ async def trigger_scheduled_task(task_id: str, request: Request):
|
||||
if task is None:
|
||||
raise HTTPException(status_code=404, detail="Scheduled task not found")
|
||||
result = await service.dispatch_task(task, now=datetime.now(UTC), trigger="manual")
|
||||
if result["outcome"] == "not_found":
|
||||
raise HTTPException(status_code=404, detail=result["error"] or "Scheduled task not found")
|
||||
if result["outcome"] == "conflict":
|
||||
raise HTTPException(status_code=409, detail=result["error"] or "Scheduled task trigger conflicted with an active run")
|
||||
if result["outcome"] == "failed":
|
||||
@ -276,10 +326,20 @@ async def delete_scheduled_task(task_id: str, request: Request):
|
||||
user = await get_optional_user_from_request(request)
|
||||
if user is None:
|
||||
raise HTTPException(status_code=401, detail="Authentication required")
|
||||
deleted = await repo.delete(task_id, user_id=str(user.id))
|
||||
if not deleted:
|
||||
result = await repo.delete_with_queue_cancellation(
|
||||
task_id,
|
||||
user_id=str(user.id),
|
||||
error="scheduled task was deleted while queued",
|
||||
now=datetime.now(UTC),
|
||||
)
|
||||
if result == "not_found":
|
||||
raise HTTPException(status_code=404, detail="Scheduled task not found")
|
||||
return {"id": task_id, "deleted": deleted}
|
||||
if result == "executing":
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Scheduled task is already launching or running; retry after the active execution finishes",
|
||||
)
|
||||
return {"id": task_id, "deleted": True}
|
||||
|
||||
|
||||
@router.get("/scheduled-tasks/{task_id}/runs")
|
||||
|
||||
@ -1478,7 +1478,14 @@ async def launch_scheduled_thread_run(
|
||||
if_not_exists="create",
|
||||
feedback_keys=None,
|
||||
)
|
||||
record = await start_run(body, thread_id, request)
|
||||
scheduled_task_run_id = (metadata or {}).get("scheduled_task_run_id")
|
||||
idempotency_key = f"scheduled-task:{scheduled_task_run_id}" if isinstance(scheduled_task_run_id, str) else None
|
||||
record = await start_run(
|
||||
body,
|
||||
thread_id,
|
||||
request,
|
||||
idempotency_key=idempotency_key,
|
||||
)
|
||||
return {"run_id": record.run_id, "thread_id": record.thread_id}
|
||||
|
||||
|
||||
|
||||
@ -4,24 +4,24 @@ import asyncio
|
||||
import logging
|
||||
import socket
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any, Literal
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from deerflow.persistence.scheduled_task_runs import ActiveScheduledRunConflict
|
||||
from deerflow.persistence.scheduled_task_runs import ActiveScheduledRunConflict, ScheduledTaskAdmissionRejected
|
||||
from deerflow.runtime import ConflictError, RunRecord
|
||||
from deerflow.scheduler.schedules import next_run_at
|
||||
from deerflow.utils.thread_id import validate_thread_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Shared so the has_active_runs fast path and the unique-index race path return
|
||||
# byte-identical outcomes for the same "task already has an active run" condition.
|
||||
# Shared so the active-row fast path and the atomic-admission conflict path
|
||||
# return byte-identical outcomes for the same active-occurrence 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"
|
||||
_QUEUE_TIMEOUT_ERROR = "scheduled task queue wait timeout exceeded"
|
||||
|
||||
|
||||
class ScheduledTaskService:
|
||||
@ -34,6 +34,7 @@ class ScheduledTaskService:
|
||||
poll_interval_seconds: int,
|
||||
lease_seconds: int,
|
||||
max_concurrent_runs: int,
|
||||
queue_timeout_seconds: int = 3600,
|
||||
multi_instance: bool = False,
|
||||
run_lease_grace_seconds: int = 10,
|
||||
) -> None:
|
||||
@ -43,6 +44,7 @@ class ScheduledTaskService:
|
||||
self._poll_interval_seconds = poll_interval_seconds
|
||||
self._lease_seconds = lease_seconds
|
||||
self._max_concurrent_runs = max_concurrent_runs
|
||||
self._queue_timeout_seconds = queue_timeout_seconds
|
||||
self._multi_instance = multi_instance
|
||||
self._run_lease_grace_seconds = run_lease_grace_seconds
|
||||
self._lease_owner = f"{socket.gethostname()}:{uuid.uuid4().hex}"
|
||||
@ -56,27 +58,22 @@ class ScheduledTaskService:
|
||||
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(
|
||||
await self._task_run_repo.recover_expired_launch_claims(
|
||||
error=_LEASE_RECOVERY_ERROR,
|
||||
now=now,
|
||||
lease_owner=self._lease_owner,
|
||||
lease_seconds=self._lease_seconds,
|
||||
limit=budget,
|
||||
)
|
||||
await self._expire_waiting_runs(now=now)
|
||||
await self._drain_queue(now=now)
|
||||
# Admission and execution capacity are separate. Due occurrences are
|
||||
# persisted even when all execution slots are busy; claim_queued_run()
|
||||
# applies the global launch budget under the database lock.
|
||||
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,
|
||||
)
|
||||
for task in claimed:
|
||||
await self.dispatch_task(task, now=now, trigger="scheduled")
|
||||
|
||||
@ -110,14 +107,6 @@ class ScheduledTaskService:
|
||||
return "paused"
|
||||
return "enabled"
|
||||
|
||||
@staticmethod
|
||||
def _task_status_for_skip(task: dict[str, Any]) -> str:
|
||||
if task["schedule_type"] == "once":
|
||||
# The single occurrence was lost to an overlapping run; "completed"
|
||||
# would claim an execution that never happened.
|
||||
return "failed"
|
||||
return "enabled"
|
||||
|
||||
async def dispatch_task(
|
||||
self,
|
||||
task: dict[str, Any],
|
||||
@ -163,55 +152,11 @@ class ScheduledTaskService:
|
||||
"thread_id": execution_thread_id,
|
||||
"error": str(exc),
|
||||
}
|
||||
# "skip" must hold for fresh-thread runs too, where every run gets a new
|
||||
# thread and the same-thread multitask ConflictError below can never
|
||||
# fire. Checked before creating this dispatch's own run row so the row
|
||||
# does not count itself as the active run. A manual trigger against an
|
||||
# active run is rejected outright (409 at the router) instead of being
|
||||
# recorded as a skipped occurrence — nothing was scheduled to happen.
|
||||
#
|
||||
# This has_active_runs check is a non-atomic fast path: it runs in its
|
||||
# own session and is separated from the create() below by await points,
|
||||
# so two concurrent dispatches (double-click / client retry / a manual
|
||||
# trigger racing the poller) can both observe no active run. The DB is
|
||||
# the atomic arbiter — the partial unique index uq_scheduled_task_run_active
|
||||
# rejects the second active insert, surfaced as ActiveScheduledRunConflict
|
||||
# and collapsed to the SAME outcome as this fast path just below.
|
||||
overlap_skip = task.get("overlap_policy", "skip") == "skip"
|
||||
if overlap_skip and await self._task_run_repo.has_active_runs(task["id"]):
|
||||
if trigger == "manual":
|
||||
return self._active_run_conflict_result(execution_thread_id)
|
||||
return await self._record_scheduled_skip(task, thread_id=execution_thread_id, now=now, trigger=trigger)
|
||||
|
||||
# Global concurrent-run budget check for manual triggers. The poller
|
||||
# enforces max_concurrent_runs via count_active_runs() before
|
||||
# claim_due_tasks(); a manual trigger bypasses that path and must apply
|
||||
# the same cap so it cannot push the active count above the limit.
|
||||
# Like the poller's count this is a non-atomic fast path; the partial
|
||||
# unique index uq_scheduled_task_run_active is the atomic arbiter that
|
||||
# rejects a second active insert for the *same task*, but there is no
|
||||
# DB-level constraint that caps the global count, so we treat this as a
|
||||
# best-effort guard consistent with how the poller enforces the budget.
|
||||
if trigger == "manual" and self._max_concurrent_runs > 0:
|
||||
active = await self._task_run_repo.count_active_runs()
|
||||
if active >= self._max_concurrent_runs:
|
||||
return {
|
||||
"outcome": "conflict",
|
||||
"task_run_id": None,
|
||||
"run_id": None,
|
||||
"thread_id": execution_thread_id,
|
||||
"error": "global concurrent-run limit reached",
|
||||
}
|
||||
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
|
||||
active = await self._task_run_repo.get_active_run(task["id"])
|
||||
if active is not None:
|
||||
if trigger == "scheduled":
|
||||
await self._release_admission_lease(task, trigger=trigger)
|
||||
return self._existing_active_result(active, execution_thread_id, trigger=trigger)
|
||||
|
||||
task_run_id = f"task-run-{uuid.uuid4().hex}"
|
||||
try:
|
||||
@ -222,28 +167,81 @@ class ScheduledTaskService:
|
||||
scheduled_for=now,
|
||||
trigger=trigger,
|
||||
status="queued",
|
||||
coordinate_with_task=True,
|
||||
expected_task_user_id=task.get("user_id"),
|
||||
expected_task_status=task.get("status") if trigger == "manual" else None,
|
||||
expected_task_updated_at=task.get("updated_at") if trigger == "manual" else None,
|
||||
expected_task_lease_owner=self._lease_owner if trigger == "scheduled" else None,
|
||||
release_task_lease_status="enabled" if trigger == "scheduled" else None,
|
||||
)
|
||||
except ActiveScheduledRunConflict:
|
||||
# Lost the create race for the task's single active slot: a
|
||||
# concurrent dispatch passed the same fast-path check and inserted
|
||||
# its active row first. Identical outcome to the fast path above.
|
||||
if trigger == "manual":
|
||||
active = await self._task_run_repo.get_active_run(task["id"])
|
||||
if trigger == "scheduled":
|
||||
await self._release_admission_lease(task, trigger=trigger)
|
||||
if active is None:
|
||||
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)
|
||||
return self._existing_active_result(active, execution_thread_id, trigger=trigger)
|
||||
except ScheduledTaskAdmissionRejected as exc:
|
||||
if exc.reason == "not_found":
|
||||
return {
|
||||
"outcome": "not_found",
|
||||
"task_run_id": None,
|
||||
"run_id": None,
|
||||
"thread_id": execution_thread_id,
|
||||
"error": "scheduled task no longer exists",
|
||||
}
|
||||
return {
|
||||
"outcome": "conflict",
|
||||
"task_run_id": None,
|
||||
"run_id": None,
|
||||
"thread_id": execution_thread_id,
|
||||
"error": "scheduled task changed before trigger admission",
|
||||
}
|
||||
|
||||
# Scheduled admission inserted the queue row and released its parent
|
||||
# lease in one transaction. Manual admission verified that this task
|
||||
# snapshot was still current under the same parent lock.
|
||||
queued = {
|
||||
"id": task_run_id,
|
||||
"task_id": task["id"],
|
||||
"thread_id": execution_thread_id,
|
||||
"trigger": trigger,
|
||||
}
|
||||
return await self._attempt_queued_run(task, queued, now=now)
|
||||
|
||||
async def _release_admission_lease(self, task: dict[str, Any], *, trigger: str) -> None:
|
||||
status = "enabled" if trigger == "scheduled" else (task.get("status") or "enabled")
|
||||
await self._task_repo.release_dispatch_lease(
|
||||
task["id"],
|
||||
expected_lease_owner=self._lease_owner if trigger == "scheduled" else None,
|
||||
status=status,
|
||||
)
|
||||
|
||||
async def _attempt_queued_run(
|
||||
self,
|
||||
task: dict[str, Any],
|
||||
queued: dict[str, Any],
|
||||
*,
|
||||
now: datetime,
|
||||
) -> dict[str, Any]:
|
||||
task_run_id = queued["id"]
|
||||
execution_thread_id = queued["thread_id"]
|
||||
trigger = queued["trigger"]
|
||||
claimed = await self._task_run_repo.claim_queued_run(
|
||||
task_run_id,
|
||||
lease_owner=self._lease_owner,
|
||||
now=now,
|
||||
lease_seconds=self._lease_seconds,
|
||||
global_max_concurrent_runs=self._max_concurrent_runs,
|
||||
)
|
||||
if claimed is None:
|
||||
return self._queued_result(task_run_id, execution_thread_id)
|
||||
|
||||
# Track whether _launch_run has produced a live run. A bookkeeping
|
||||
# failure AFTER launch (the queued->running write, or the parent task
|
||||
# update) must NOT be recorded as "failed": "failed" is outside the
|
||||
# partial unique index uq_scheduled_task_run_active, so it would release
|
||||
# the task's single active slot and the next dispatch would launch a
|
||||
# duplicate run. Once launch succeeds we keep the row "running" and
|
||||
# retain the launched run_id regardless of bookkeeping errors.
|
||||
# failure after launch must retain the non-terminal slot so a later
|
||||
# poll cannot start the same occurrence twice.
|
||||
launched_run_id: str | None = None
|
||||
launched_thread_id: str | None = None
|
||||
# Flip immediately after _launch_run returns, before any further code
|
||||
# that can raise (e.g. result["run_id"] on a malformed result). The
|
||||
# retention branch keys off this flag, not `launched_run_id is not
|
||||
# None`, so a launch that succeeded but whose result-unpacking raised
|
||||
# still takes the retention path instead of the release-the-slot path.
|
||||
launch_succeeded = False
|
||||
try:
|
||||
result = await self._launch_run(
|
||||
@ -267,14 +265,11 @@ class ScheduledTaskService:
|
||||
now=now,
|
||||
)
|
||||
task_status = self._task_status_for_launch(task, trigger=trigger)
|
||||
await self._task_run_repo.update_status(
|
||||
task_run_id,
|
||||
status="running",
|
||||
await self._record_launched_run(
|
||||
task_run_id=task_run_id,
|
||||
task_id=task["id"],
|
||||
run_id=launched_run_id,
|
||||
started_at=now,
|
||||
# A fast-failing run can reach handle_run_completion before this
|
||||
# write resumes; never clobber its terminal status.
|
||||
protect_terminal=True,
|
||||
)
|
||||
await self._task_repo.update_after_launch(
|
||||
task["id"],
|
||||
@ -288,7 +283,6 @@ 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",
|
||||
@ -298,19 +292,13 @@ class ScheduledTaskService:
|
||||
"error": None,
|
||||
}
|
||||
except Exception as exc:
|
||||
if not launch_succeeded and self._is_overlap_conflict(exc) and trigger == "scheduled" and task.get("overlap_policy", "skip") == "skip":
|
||||
# Pre-launch overlap conflict (e.g. same-thread multitask): no
|
||||
# run was started, so recording a skip and releasing the slot is
|
||||
# safe. Guarded by ``not launch_succeeded`` because a run that
|
||||
# already launched must never be reclassified as a skip / failed.
|
||||
return await self._finalize_skip(
|
||||
task,
|
||||
task_run_id=task_run_id,
|
||||
thread_id=execution_thread_id,
|
||||
now=now,
|
||||
if not launch_succeeded and self._is_overlap_conflict(exc):
|
||||
await self._task_run_repo.requeue_claimed_run(
|
||||
task_run_id,
|
||||
lease_owner=self._lease_owner,
|
||||
error=str(exc),
|
||||
trigger=trigger,
|
||||
)
|
||||
return self._queued_result(task_run_id, execution_thread_id, error=str(exc))
|
||||
|
||||
next_at = next_run_at(
|
||||
task["schedule_type"],
|
||||
@ -331,12 +319,11 @@ class ScheduledTaskService:
|
||||
# the run as launched so callers know a run is in flight.
|
||||
task_status = self._task_status_for_launch(task, trigger=trigger)
|
||||
try:
|
||||
await self._task_run_repo.update_status(
|
||||
task_run_id,
|
||||
status="running",
|
||||
await self._record_launched_run(
|
||||
task_run_id=task_run_id,
|
||||
task_id=task["id"],
|
||||
run_id=launched_run_id,
|
||||
started_at=now,
|
||||
protect_terminal=True,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
@ -363,7 +350,6 @@ class ScheduledTaskService:
|
||||
last_error=None,
|
||||
increment_run_count=True,
|
||||
protect_terminal=True,
|
||||
expected_lease_owner=expected_lease_owner,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
@ -381,33 +367,58 @@ class ScheduledTaskService:
|
||||
|
||||
# _launch_run itself failed (or a step before it did): no live run
|
||||
# was created, so it is safe to release the active slot.
|
||||
task_status = self._task_status_for_failure(task, trigger=trigger)
|
||||
await self._task_run_repo.update_status(
|
||||
finalized = await self._task_run_repo.fail_launching_run(
|
||||
task_run_id,
|
||||
status="failed",
|
||||
task_id=task["id"],
|
||||
lease_owner=self._lease_owner,
|
||||
error=str(exc),
|
||||
started_at=now,
|
||||
finished_at=now,
|
||||
)
|
||||
await self._task_repo.update_after_launch(
|
||||
task["id"],
|
||||
status=task_status,
|
||||
next_run_at=next_at,
|
||||
last_run_at=now,
|
||||
last_run_id=None,
|
||||
last_thread_id=execution_thread_id,
|
||||
last_error=str(exc),
|
||||
increment_run_count=False,
|
||||
expected_lease_owner=expected_lease_owner,
|
||||
now=now,
|
||||
)
|
||||
if not finalized:
|
||||
logger.warning(
|
||||
"Scheduled task-run %s lost its launch claim before failure bookkeeping; leaving recovery-owned state unchanged",
|
||||
task_run_id,
|
||||
)
|
||||
return self._queued_result(task_run_id, execution_thread_id, error=str(exc))
|
||||
return {
|
||||
"outcome": "conflict" if self._is_overlap_conflict(exc) else "failed",
|
||||
"outcome": "failed",
|
||||
"task_run_id": task_run_id,
|
||||
"run_id": None,
|
||||
"thread_id": execution_thread_id,
|
||||
"error": str(exc),
|
||||
}
|
||||
|
||||
async def _record_launched_run(
|
||||
self,
|
||||
*,
|
||||
task_run_id: str,
|
||||
task_id: str,
|
||||
run_id: str,
|
||||
started_at: datetime,
|
||||
) -> None:
|
||||
updated = await self._task_run_repo.update_status(
|
||||
task_run_id,
|
||||
status="running",
|
||||
run_id=run_id,
|
||||
started_at=started_at,
|
||||
protect_terminal=True,
|
||||
expected_lease_owner=self._lease_owner,
|
||||
)
|
||||
if updated:
|
||||
return
|
||||
reconciled = await self._task_run_repo.reconcile_launched_run(
|
||||
task_run_id,
|
||||
task_id=task_id,
|
||||
run_id=run_id,
|
||||
started_at=started_at,
|
||||
)
|
||||
if not reconciled:
|
||||
logger.error(
|
||||
"Scheduled task-run %s launched durable run %s but could not restore its occurrence association",
|
||||
task_run_id,
|
||||
run_id,
|
||||
)
|
||||
|
||||
def _active_run_conflict_result(self, thread_id: str) -> dict[str, Any]:
|
||||
"""Manual-trigger response when the task already has an active run.
|
||||
|
||||
@ -422,82 +433,66 @@ class ScheduledTaskService:
|
||||
"error": _ACTIVE_RUN_CONFLICT_ERROR,
|
||||
}
|
||||
|
||||
async def _record_scheduled_skip(
|
||||
def _existing_active_result(
|
||||
self,
|
||||
task: dict[str, Any],
|
||||
*,
|
||||
active: dict[str, Any],
|
||||
thread_id: str,
|
||||
now: datetime,
|
||||
*,
|
||||
trigger: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Record a skipped occurrence for a scheduled dispatch that overlapped an active run.
|
||||
if active["status"] == "queued":
|
||||
return self._queued_result(active["id"], active["thread_id"])
|
||||
return self._active_run_conflict_result(thread_id)
|
||||
|
||||
The tombstone is created directly as terminal ``"skipped"`` rather than
|
||||
the transient ``"queued"`` the launch path uses: a queued row is active
|
||||
and would itself trip ``uq_scheduled_task_run_active`` against the
|
||||
pre-existing run that is still holding the task's single active slot.
|
||||
``"skipped"`` is outside the index predicate, so it never conflicts.
|
||||
"""
|
||||
task_run_id = f"task-run-{uuid.uuid4().hex}"
|
||||
await self._task_run_repo.create(
|
||||
run_record_id=task_run_id,
|
||||
task_id=task["id"],
|
||||
thread_id=thread_id,
|
||||
scheduled_for=now,
|
||||
trigger=trigger,
|
||||
status="skipped",
|
||||
)
|
||||
return await self._finalize_skip(
|
||||
task,
|
||||
task_run_id=task_run_id,
|
||||
thread_id=thread_id,
|
||||
now=now,
|
||||
error=_SKIP_ACTIVE_RUN_ERROR,
|
||||
trigger=trigger,
|
||||
)
|
||||
|
||||
async def _finalize_skip(
|
||||
self,
|
||||
task: dict[str, Any],
|
||||
*,
|
||||
@staticmethod
|
||||
def _queued_result(
|
||||
task_run_id: str,
|
||||
thread_id: str,
|
||||
now: datetime,
|
||||
error: str,
|
||||
trigger: str,
|
||||
*,
|
||||
error: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
next_at = next_run_at(
|
||||
task["schedule_type"],
|
||||
task["schedule_spec"],
|
||||
task["timezone"],
|
||||
now=now,
|
||||
)
|
||||
await self._task_run_repo.update_status(
|
||||
task_run_id,
|
||||
status="skipped",
|
||||
error=error,
|
||||
started_at=now,
|
||||
finished_at=now,
|
||||
)
|
||||
await self._task_repo.update_after_launch(
|
||||
task["id"],
|
||||
status=self._task_status_for_skip(task),
|
||||
next_run_at=next_at,
|
||||
last_run_at=task.get("last_run_at"),
|
||||
last_run_id=task.get("last_run_id"),
|
||||
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",
|
||||
"outcome": "queued",
|
||||
"task_run_id": task_run_id,
|
||||
"run_id": None,
|
||||
"thread_id": thread_id,
|
||||
"error": error,
|
||||
}
|
||||
|
||||
async def _drain_queue(self, *, now: datetime) -> None:
|
||||
queued_rows = await self._task_run_repo.list_queued_runs(limit=max(16, self._max_concurrent_runs * 4))
|
||||
for queued in queued_rows:
|
||||
await self._task_repo.release_queued_admission_lease(queued["task_id"])
|
||||
task = await self._task_repo.get_internal(queued["task_id"])
|
||||
if task is None:
|
||||
await self._task_run_repo.update_status(
|
||||
queued["id"],
|
||||
status="interrupted",
|
||||
error="scheduled task was deleted while queued",
|
||||
finished_at=now,
|
||||
)
|
||||
continue
|
||||
# Pausing suppresses automatic occurrences, but a manual trigger is
|
||||
# an explicit request and has always been allowed to run without
|
||||
# resuming the schedule. A later pause still cancels an already
|
||||
# queued manual row atomically in pause_with_queue_cancellation().
|
||||
if task.get("status") == "paused" and queued["trigger"] != "manual":
|
||||
await self._task_run_repo.update_status(
|
||||
queued["id"],
|
||||
status="interrupted",
|
||||
error="scheduled task was paused while queued",
|
||||
finished_at=now,
|
||||
)
|
||||
continue
|
||||
await self._attempt_queued_run(task, queued, now=now)
|
||||
|
||||
async def _expire_waiting_runs(self, *, now: datetime) -> None:
|
||||
await self._task_run_repo.expire_queued_runs(
|
||||
created_before=now - timedelta(seconds=self._queue_timeout_seconds),
|
||||
error=_QUEUE_TIMEOUT_ERROR,
|
||||
now=now,
|
||||
)
|
||||
|
||||
async def handle_run_completion(self, record: RunRecord) -> None:
|
||||
metadata = record.metadata or {}
|
||||
task_id = metadata.get("scheduled_task_id")
|
||||
|
||||
@ -241,6 +241,7 @@ scheduler:
|
||||
poll_interval_seconds: 5
|
||||
lease_seconds: 120
|
||||
max_concurrent_runs: 3
|
||||
queue_timeout_seconds: 3600
|
||||
min_once_delay_seconds: 60
|
||||
recursion_limit: 1000
|
||||
```
|
||||
@ -249,12 +250,16 @@ Notes:
|
||||
|
||||
- `enabled: false` keeps background polling off by default.
|
||||
- `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.
|
||||
- `max_concurrent_runs` is a shared global execution cap in multi-instance mode. Waiting `queued` rows do not consume capacity; an atomic `queued` → `launching` claim counts `launching`/`running` rows under a Postgres advisory lock so concurrent Pods cannot exceed the cap.
|
||||
- `queue_timeout_seconds` limits how long a persisted occurrence may wait for capacity or a reused thread to become available. Expired occurrences are marked `failed`; queued rows otherwise survive Gateway restarts.
|
||||
- A task definition is immutable while an occurrence is `queued`, `launching`, or `running`. This prevents a durable occurrence from mixing its admitted thread with a later prompt or schedule edit. Transitioning a task to paused or deleting it cancels a waiting row; PATCH and resume return a conflict until the active occurrence finishes or is cancelled.
|
||||
- A manual trigger remains explicit even while the recurring schedule is paused: it may wait in the durable queue and run later, while the task itself stays paused. Transitioning an enabled task to paused still cancels its waiting occurrence atomically.
|
||||
- Queue admission, PATCH/resume, pause, and delete serialize on the parent task row. Per-thread FIFO spans all active states, so an older `launching` or `running` occurrence blocks a newer queued occurrence on the same reused thread as well as an older `queued` occurrence.
|
||||
- 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.
|
||||
- `recursion_limit` is the LangGraph super-step cap for scheduler-launched runs (default 1000, matching the web UI's interactive budget). Values above `max_recursion_limit` (default 1000) are clamped. This field is read at dispatch, so a YAML edit applies to the next scheduled run without a Gateway restart.
|
||||
- Poller fields (`enabled`, `multi_instance`, `poll_interval_seconds`, `lease_seconds`, `max_concurrent_runs`, `min_once_delay_seconds`) are restart-required; edits need a Gateway restart.
|
||||
- Poller fields (`enabled`, `multi_instance`, `poll_interval_seconds`, `lease_seconds`, `max_concurrent_runs`, `queue_timeout_seconds`, `min_once_delay_seconds`) 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:** in multi-instance mode, `max_concurrent_runs` is cluster-wide rather than per Pod and counts `launching`/`running` occurrences. Waiting `queued` rows remain outside the execution cap; capacity 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.
|
||||
|
||||
@ -68,7 +68,7 @@ STARTUP_ONLY_FIELDS: dict[str, str] = {
|
||||
),
|
||||
"scheduler": (
|
||||
"ScheduledTaskService is constructed and started once during Gateway lifespan startup; enabled, poll_interval_seconds, lease_seconds, "
|
||||
"max_concurrent_runs, and multi_instance are captured into the service instance and the background poller task is not rebuilt on config.yaml edits. "
|
||||
"max_concurrent_runs, queue_timeout_seconds, 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. "
|
||||
"scheduler.recursion_limit is not captured there: launch_scheduled_thread_run reads it from get_app_config() on each dispatch, so a YAML edit applies to the next scheduled run without a Gateway restart."
|
||||
),
|
||||
|
||||
@ -7,6 +7,7 @@ class SchedulerConfig(BaseModel):
|
||||
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)
|
||||
queue_timeout_seconds: int = Field(default=3600, ge=60, le=604800)
|
||||
min_once_delay_seconds: int = Field(default=60, ge=1, le=86400)
|
||||
recursion_limit: int = Field(
|
||||
default=1000,
|
||||
|
||||
@ -84,6 +84,7 @@ on installs that never enabled it. The convention is:
|
||||
- `migrations/versions/0012_mcp_task_results.py` — adds bounded result preview/truncation/artifact fields for ordinary task drivers
|
||||
- `migrations/versions/0013_mcp_task_notifications.py` — adds durable Agent-run notification snapshots, delivery leases, idempotency fields, and the separate bounded-retry attempt counter
|
||||
- `migrations/versions/0014_managed_subagents.py` — creates the deployment-level managed Subagent catalog table
|
||||
- `migrations/versions/0015_scheduled_task_enqueue.py` — interrupts legacy transient queued rows, adds durable scheduled-run launch leases and attempt counts, expands the one-active-occurrence index to `queued`/`launching`/`running`, and migrates the overlap policy from `skip` to `enqueue`; chains after `0014_managed_subagents`
|
||||
- `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch decision + locking
|
||||
- `extensions/loader.py::load_extensions` — registers each spec's `table_prefix` with `register_extension_table_prefix()`
|
||||
- Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter, including extension-owned tables), `tests/test_extension_loader.py::TestTablePrefixRegistration` (spec-to-filter wiring), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor), `tests/test_migration_0004_run_ownership_dedupe.py` + `tests/test_migration_0007_scheduled_run_active_dedupe.py` (dedupe-before-unique-index pre-steps)
|
||||
|
||||
@ -0,0 +1,72 @@
|
||||
"""Durable scheduled-task enqueue state.
|
||||
|
||||
Revision ID: 0015_scheduled_task_enqueue
|
||||
Revises: 0014_managed_subagents
|
||||
Create Date: 2026-08-20
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision: str = "0015_scheduled_task_enqueue"
|
||||
down_revision: str | Sequence[str] | None = "0014_managed_subagents"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _replace_active_index(*, statuses: str) -> None:
|
||||
bind = op.get_bind()
|
||||
existing = {index["name"] for index in sa.inspect(bind).get_indexes("scheduled_task_runs")}
|
||||
with op.batch_alter_table("scheduled_task_runs", schema=None) as batch_op:
|
||||
if "uq_scheduled_task_run_active" in existing:
|
||||
batch_op.drop_index("uq_scheduled_task_run_active")
|
||||
predicate = sa.text(f"status IN ({statuses})")
|
||||
batch_op.create_index(
|
||||
"uq_scheduled_task_run_active",
|
||||
["task_id"],
|
||||
unique=True,
|
||||
sqlite_where=predicate,
|
||||
postgresql_where=predicate,
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
from deerflow.persistence.migrations._helpers import safe_add_column
|
||||
|
||||
# Before this revision, ``queued`` was a transient pre-launch marker and
|
||||
# startup interrupted crash leftovers. Durable enqueue gives the same
|
||||
# value a new meaning, so preserve the old restart behavior at the upgrade
|
||||
# boundary instead of launching an occurrence that may already have run.
|
||||
op.execute(
|
||||
sa.text("UPDATE scheduled_task_runs SET status = 'interrupted', error = 'interrupted: gateway upgraded before the queued run reached a terminal state', finished_at = COALESCE(finished_at, CURRENT_TIMESTAMP) WHERE status = 'queued'")
|
||||
)
|
||||
safe_add_column(
|
||||
"scheduled_task_runs",
|
||||
sa.Column("lease_owner", sa.String(length=128), nullable=True),
|
||||
)
|
||||
safe_add_column(
|
||||
"scheduled_task_runs",
|
||||
sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True),
|
||||
)
|
||||
safe_add_column(
|
||||
"scheduled_task_runs",
|
||||
sa.Column("attempt_count", sa.Integer(), nullable=False, server_default="0"),
|
||||
)
|
||||
_replace_active_index(statuses="'queued', 'launching', 'running'")
|
||||
op.execute(sa.text("UPDATE scheduled_tasks SET overlap_policy = 'enqueue' WHERE overlap_policy = 'skip'"))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
from deerflow.persistence.migrations._helpers import safe_drop_column
|
||||
|
||||
# A deployment must not leave a status unknown to the old scheduler.
|
||||
op.execute(sa.text("UPDATE scheduled_task_runs SET status = 'queued', lease_owner = NULL, lease_expires_at = NULL WHERE status = 'launching'"))
|
||||
op.execute(sa.text("UPDATE scheduled_tasks SET overlap_policy = 'skip' WHERE overlap_policy = 'enqueue'"))
|
||||
_replace_active_index(statuses="'queued', 'running'")
|
||||
safe_drop_column("scheduled_task_runs", "attempt_count")
|
||||
safe_drop_column("scheduled_task_runs", "lease_expires_at")
|
||||
safe_drop_column("scheduled_task_runs", "lease_owner")
|
||||
@ -1,4 +1,9 @@
|
||||
from .model import ScheduledTaskRunRow
|
||||
from .sql import ActiveScheduledRunConflict, ScheduledTaskRunRepository
|
||||
from .sql import ActiveScheduledRunConflict, ScheduledTaskAdmissionRejected, ScheduledTaskRunRepository
|
||||
|
||||
__all__ = ["ActiveScheduledRunConflict", "ScheduledTaskRunRow", "ScheduledTaskRunRepository"]
|
||||
__all__ = [
|
||||
"ActiveScheduledRunConflict",
|
||||
"ScheduledTaskAdmissionRejected",
|
||||
"ScheduledTaskRunRow",
|
||||
"ScheduledTaskRunRepository",
|
||||
]
|
||||
|
||||
@ -2,7 +2,7 @@ from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from sqlalchemy import DateTime, Index, String, Text, text
|
||||
from sqlalchemy import DateTime, Index, Integer, String, Text, text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from deerflow.persistence.base import Base
|
||||
@ -19,27 +19,23 @@ class ScheduledTaskRunRow(Base):
|
||||
trigger: Mapped[str] = mapped_column(String(16))
|
||||
status: Mapped[str] = mapped_column(String(16), index=True)
|
||||
error: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
lease_owner: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||||
lease_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
attempt_count: Mapped[int] = mapped_column(Integer, default=0, server_default="0")
|
||||
started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
|
||||
|
||||
__table_args__ = (
|
||||
# At most one active (queued/running) run per task. This is the atomic
|
||||
# arbiter for the ``dispatch_task`` skip policy: the non-atomic
|
||||
# ``has_active_runs`` check-then-create is a fast path, but two
|
||||
# concurrent dispatches (double-click / client retry / manual trigger
|
||||
# racing the poller) can both pass it, so the DB must reject the second
|
||||
# active insert. Sibling of the ``runs`` table's ``uq_runs_thread_active``
|
||||
# At most one non-terminal (queued/launching/running) occurrence per
|
||||
# task. Queued occurrences are deliberately durable; ``launching`` is
|
||||
# a short lease-fenced claim used so multiple gateway instances cannot
|
||||
# launch the same row. Sibling of the ``runs`` table's
|
||||
# ``uq_runs_thread_active``
|
||||
# (PR #4003); that one keys on ``thread_id`` and does not cover the
|
||||
# default ``fresh_thread_per_run`` context (every dispatch gets a new
|
||||
# thread), which is why the scheduled-task run row needs its own guard.
|
||||
#
|
||||
# Condition is status-only, not ``overlap_policy``: the policy is fixed
|
||||
# to "skip" in the MVP, so a status-only predicate enforces the current
|
||||
# invariant without denormalizing ``overlap_policy`` onto the run row
|
||||
# for an unimplemented non-skip policy. If a non-skip policy is added
|
||||
# this must become conditional (e.g. ``... AND overlap_policy = 'skip'``).
|
||||
#
|
||||
# Must live in ORM ``__table_args__`` (not just the migration) because
|
||||
# the empty-DB bootstrap path runs ``create_all`` + ``stamp head`` and
|
||||
# never executes the migration that also defines this index.
|
||||
@ -47,7 +43,7 @@ class ScheduledTaskRunRow(Base):
|
||||
"uq_scheduled_task_run_active",
|
||||
"task_id",
|
||||
unique=True,
|
||||
sqlite_where=text("status IN ('queued', 'running')"),
|
||||
postgresql_where=text("status IN ('queued', 'running')"),
|
||||
sqlite_where=text("status IN ('queued', 'launching', 'running')"),
|
||||
postgresql_where=text("status IN ('queued', 'launching', 'running')"),
|
||||
),
|
||||
)
|
||||
|
||||
@ -3,18 +3,23 @@ from __future__ import annotations
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import func, select
|
||||
from sqlalchemy import and_, exists, func, or_, select, text, update
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
from sqlalchemy.orm import aliased
|
||||
|
||||
from deerflow.persistence.run import RunRepository
|
||||
from deerflow.persistence.run.model import RunRow
|
||||
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
|
||||
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow
|
||||
from deerflow.scheduler.schedules import next_run_at as compute_next_run_at
|
||||
from deerflow.utils.time import coerce_iso
|
||||
|
||||
TERMINAL_RUN_STATUSES: frozenset[str] = frozenset({"success", "failed", "skipped", "interrupted"})
|
||||
ACTIVE_RUN_STATUSES: tuple[str, ...] = ("queued", "running")
|
||||
QUEUED_RUN_STATUSES: tuple[str, ...] = ("queued",)
|
||||
EXECUTING_RUN_STATUSES: tuple[str, ...] = ("launching", "running")
|
||||
ACTIVE_RUN_STATUSES: tuple[str, ...] = (*QUEUED_RUN_STATUSES, *EXECUTING_RUN_STATUSES)
|
||||
_SCHEDULER_BUDGET_LOCK_KEY = 4694001
|
||||
|
||||
|
||||
def _lease_is_alive(lease_expires_at: datetime | None, *, now: datetime, grace_seconds: int) -> bool:
|
||||
@ -26,18 +31,12 @@ def _lease_is_alive(lease_expires_at: datetime | None, *, now: datetime, grace_s
|
||||
|
||||
|
||||
class ActiveScheduledRunConflict(Exception):
|
||||
"""A concurrent dispatch already holds the task's single active-run slot.
|
||||
"""A concurrent dispatch already holds the task's active-occurrence slot.
|
||||
|
||||
Raised by :meth:`ScheduledTaskRunRepository.create` when inserting an
|
||||
active (queued/running) run row would violate the partial unique index
|
||||
``uq_scheduled_task_run_active`` (at most one active run per ``task_id``).
|
||||
This is the atomic counterpart to the non-atomic ``has_active_runs`` check
|
||||
in ``ScheduledTaskService.dispatch_task``: two dispatches can both pass that
|
||||
check, but only one can insert the active row — the loser lands here.
|
||||
|
||||
Translating the SQLAlchemy ``IntegrityError`` into a domain exception at
|
||||
the repository boundary keeps the service layer free of ``sqlalchemy.exc``
|
||||
coupling (mirrors ``deerflow.runtime.ConflictError`` for the runs table).
|
||||
Coordinated admission serializes on the parent task before inserting the
|
||||
queue row. The partial unique index remains the database backstop for
|
||||
direct repository callers and legacy interleavings. Both paths surface the
|
||||
same domain exception without coupling the service to SQLAlchemy errors.
|
||||
"""
|
||||
|
||||
def __init__(self, task_id: str) -> None:
|
||||
@ -45,6 +44,15 @@ class ActiveScheduledRunConflict(Exception):
|
||||
super().__init__(f"scheduled task {task_id!r} already has an active run")
|
||||
|
||||
|
||||
class ScheduledTaskAdmissionRejected(Exception):
|
||||
"""The parent task changed or disappeared before queue admission."""
|
||||
|
||||
def __init__(self, task_id: str, *, reason: str) -> None:
|
||||
self.task_id = task_id
|
||||
self.reason = reason
|
||||
super().__init__(f"scheduled task {task_id!r} admission rejected: {reason}")
|
||||
|
||||
|
||||
class ScheduledTaskRunRepository:
|
||||
def __init__(
|
||||
self,
|
||||
@ -58,11 +66,77 @@ class ScheduledTaskRunRepository:
|
||||
@staticmethod
|
||||
def _row_to_dict(row: ScheduledTaskRunRow) -> dict[str, Any]:
|
||||
data = row.to_dict()
|
||||
for key in ("scheduled_for", "started_at", "finished_at", "created_at"):
|
||||
for key in (
|
||||
"scheduled_for",
|
||||
"lease_expires_at",
|
||||
"started_at",
|
||||
"finished_at",
|
||||
"created_at",
|
||||
):
|
||||
if data.get(key) is not None:
|
||||
data[key] = coerce_iso(data[key])
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
async def _lock_task(session: AsyncSession, task_id: str) -> ScheduledTaskRow | None:
|
||||
# SQLite ignores SELECT ... FOR UPDATE. Touch the parent first so its
|
||||
# single-writer lock provides the same serialization point used by
|
||||
# Postgres row locking for admission, mutation, pause, and delete.
|
||||
if session.get_bind().dialect.name == "sqlite":
|
||||
await session.execute(update(ScheduledTaskRow).where(ScheduledTaskRow.id == task_id).values(updated_at=ScheduledTaskRow.updated_at))
|
||||
return await session.get(ScheduledTaskRow, task_id, with_for_update=True)
|
||||
|
||||
@staticmethod
|
||||
def _associate_scheduled_run(
|
||||
row: ScheduledTaskRunRow,
|
||||
candidate: RunRow,
|
||||
) -> None:
|
||||
"""Fill launch bookkeeping that the durable run already proves."""
|
||||
row.run_id = candidate.run_id
|
||||
if row.started_at is None:
|
||||
row.started_at = candidate.created_at
|
||||
|
||||
@staticmethod
|
||||
def _associate_task_with_run(
|
||||
task: ScheduledTaskRow | None,
|
||||
row: ScheduledTaskRunRow,
|
||||
candidate: RunRow,
|
||||
) -> None:
|
||||
"""Repair the parent update if launch committed before bookkeeping."""
|
||||
if task is None or task.last_run_id == candidate.run_id:
|
||||
return
|
||||
launched_at = candidate.created_at
|
||||
if launched_at.tzinfo is None:
|
||||
launched_at = launched_at.replace(tzinfo=UTC)
|
||||
task.last_run_at = launched_at
|
||||
task.last_run_id = candidate.run_id
|
||||
task.last_thread_id = row.thread_id
|
||||
task.next_run_at = compute_next_run_at(
|
||||
task.schedule_type,
|
||||
task.schedule_spec,
|
||||
task.timezone,
|
||||
now=launched_at,
|
||||
)
|
||||
task.run_count += 1
|
||||
task.lease_owner = None
|
||||
task.lease_expires_at = None
|
||||
if task.schedule_type == "once":
|
||||
if candidate.status == "success":
|
||||
task.status = "completed"
|
||||
task.last_error = None
|
||||
elif candidate.status in {"error", "timeout"}:
|
||||
task.status = "failed"
|
||||
task.last_error = candidate.error
|
||||
elif candidate.status == "interrupted":
|
||||
task.status = "cancelled"
|
||||
task.last_error = candidate.error
|
||||
else:
|
||||
task.status = "running"
|
||||
task.last_error = None
|
||||
elif not (row.trigger == "manual" and task.status == "paused"):
|
||||
task.status = "enabled"
|
||||
task.last_error = candidate.error if candidate.status in {"error", "timeout", "interrupted"} else None
|
||||
|
||||
async def create(
|
||||
self,
|
||||
*,
|
||||
@ -72,6 +146,12 @@ class ScheduledTaskRunRepository:
|
||||
scheduled_for: datetime,
|
||||
trigger: str,
|
||||
status: str,
|
||||
coordinate_with_task: bool = False,
|
||||
expected_task_user_id: str | None = None,
|
||||
expected_task_status: str | None = None,
|
||||
expected_task_updated_at: datetime | str | None = None,
|
||||
expected_task_lease_owner: str | None = None,
|
||||
release_task_lease_status: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
row = ScheduledTaskRunRow(
|
||||
id=run_record_id,
|
||||
@ -83,7 +163,40 @@ class ScheduledTaskRunRepository:
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
async with self._sf() as session:
|
||||
task: ScheduledTaskRow | None = None
|
||||
if coordinate_with_task:
|
||||
task = await self._lock_task(session, task_id)
|
||||
if task is None or (expected_task_user_id is not None and task.user_id != expected_task_user_id):
|
||||
await session.rollback()
|
||||
raise ScheduledTaskAdmissionRejected(task_id, reason="not_found")
|
||||
if expected_task_lease_owner is not None:
|
||||
if task.lease_owner != expected_task_lease_owner:
|
||||
await session.rollback()
|
||||
raise ScheduledTaskAdmissionRejected(task_id, reason="stale")
|
||||
else:
|
||||
if expected_task_status is not None and task.status != expected_task_status:
|
||||
await session.rollback()
|
||||
raise ScheduledTaskAdmissionRejected(task_id, reason="stale")
|
||||
if expected_task_updated_at is not None and coerce_iso(task.updated_at) != coerce_iso(expected_task_updated_at):
|
||||
await session.rollback()
|
||||
raise ScheduledTaskAdmissionRejected(task_id, reason="stale")
|
||||
active_status = await session.scalar(
|
||||
select(ScheduledTaskRunRow.status)
|
||||
.where(
|
||||
ScheduledTaskRunRow.task_id == task_id,
|
||||
ScheduledTaskRunRow.status.in_(ACTIVE_RUN_STATUSES),
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
if active_status is not None:
|
||||
await session.rollback()
|
||||
raise ActiveScheduledRunConflict(task_id)
|
||||
session.add(row)
|
||||
if task is not None and release_task_lease_status is not None:
|
||||
task.status = release_task_lease_status
|
||||
task.lease_owner = None
|
||||
task.lease_expires_at = None
|
||||
task.updated_at = datetime.now(UTC)
|
||||
try:
|
||||
await session.commit()
|
||||
except IntegrityError:
|
||||
@ -115,12 +228,370 @@ class ScheduledTaskRunRepository:
|
||||
return [self._row_to_dict(row) for row in result.scalars()]
|
||||
|
||||
async def count_active_runs(self) -> int:
|
||||
"""Global count of queued/running rows, used to bound cross-task concurrency."""
|
||||
stmt = select(func.count()).select_from(ScheduledTaskRunRow).where(ScheduledTaskRunRow.status.in_(ACTIVE_RUN_STATUSES))
|
||||
"""Count launch claims and live runs; waiting rows do not consume slots."""
|
||||
stmt = select(func.count()).select_from(ScheduledTaskRunRow).where(ScheduledTaskRunRow.status.in_(EXECUTING_RUN_STATUSES))
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(stmt)
|
||||
return int(result.scalar() or 0)
|
||||
|
||||
async def list_queued_runs(self, *, limit: int) -> list[dict[str, Any]]:
|
||||
older = aliased(ScheduledTaskRunRow)
|
||||
older_same_thread = exists(
|
||||
select(older.id).where(
|
||||
older.thread_id == ScheduledTaskRunRow.thread_id,
|
||||
older.status.in_(ACTIVE_RUN_STATUSES),
|
||||
or_(
|
||||
older.created_at < ScheduledTaskRunRow.created_at,
|
||||
and_(
|
||||
older.created_at == ScheduledTaskRunRow.created_at,
|
||||
older.id < ScheduledTaskRunRow.id,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
stmt = (
|
||||
select(ScheduledTaskRunRow)
|
||||
.where(
|
||||
ScheduledTaskRunRow.status == "queued",
|
||||
~older_same_thread,
|
||||
)
|
||||
# Prefer rows that have had fewer launch attempts. A permanently
|
||||
# busy thread therefore cannot monopolize the bounded drain batch,
|
||||
# while created_at/id preserve FIFO order among equal attempts.
|
||||
.order_by(
|
||||
ScheduledTaskRunRow.attempt_count.asc(),
|
||||
ScheduledTaskRunRow.created_at.asc(),
|
||||
ScheduledTaskRunRow.id.asc(),
|
||||
)
|
||||
.limit(limit)
|
||||
)
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(stmt)
|
||||
return [self._row_to_dict(row) for row in result.scalars()]
|
||||
|
||||
async def get_active_run(self, task_id: str) -> dict[str, Any] | None:
|
||||
stmt = (
|
||||
select(ScheduledTaskRunRow)
|
||||
.where(
|
||||
ScheduledTaskRunRow.task_id == task_id,
|
||||
ScheduledTaskRunRow.status.in_(ACTIVE_RUN_STATUSES),
|
||||
)
|
||||
.order_by(ScheduledTaskRunRow.created_at.asc(), ScheduledTaskRunRow.id.asc())
|
||||
.limit(1)
|
||||
)
|
||||
async with self._sf() as session:
|
||||
row = (await session.execute(stmt)).scalars().first()
|
||||
return self._row_to_dict(row) if row is not None else None
|
||||
|
||||
async def claim_queued_run(
|
||||
self,
|
||||
run_record_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
now: datetime,
|
||||
lease_seconds: int,
|
||||
global_max_concurrent_runs: int,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Atomically move one waiting row into the lease-fenced launch phase."""
|
||||
async with self._sf() as session:
|
||||
if session.get_bind().dialect.name == "postgresql":
|
||||
await session.execute(
|
||||
text("SELECT pg_advisory_xact_lock(:lock_key)"),
|
||||
{"lock_key": _SCHEDULER_BUDGET_LOCK_KEY},
|
||||
)
|
||||
executing = await session.scalar(select(func.count()).select_from(ScheduledTaskRunRow).where(ScheduledTaskRunRow.status.in_(EXECUTING_RUN_STATUSES)))
|
||||
if int(executing or 0) >= global_max_concurrent_runs:
|
||||
await session.rollback()
|
||||
return None
|
||||
older = aliased(ScheduledTaskRunRow)
|
||||
older_same_thread = exists(
|
||||
select(older.id).where(
|
||||
older.thread_id == ScheduledTaskRunRow.thread_id,
|
||||
older.status.in_(ACTIVE_RUN_STATUSES),
|
||||
or_(
|
||||
older.created_at < ScheduledTaskRunRow.created_at,
|
||||
and_(
|
||||
older.created_at == ScheduledTaskRunRow.created_at,
|
||||
older.id < ScheduledTaskRunRow.id,
|
||||
),
|
||||
),
|
||||
)
|
||||
)
|
||||
result = await session.execute(
|
||||
update(ScheduledTaskRunRow)
|
||||
.where(
|
||||
ScheduledTaskRunRow.id == run_record_id,
|
||||
ScheduledTaskRunRow.status == "queued",
|
||||
~older_same_thread,
|
||||
)
|
||||
.values(
|
||||
status="launching",
|
||||
lease_owner=lease_owner,
|
||||
lease_expires_at=now + timedelta(seconds=lease_seconds),
|
||||
attempt_count=ScheduledTaskRunRow.attempt_count + 1,
|
||||
)
|
||||
)
|
||||
if result.rowcount != 1:
|
||||
await session.rollback()
|
||||
return None
|
||||
await session.commit()
|
||||
row = await session.get(ScheduledTaskRunRow, run_record_id)
|
||||
return self._row_to_dict(row) if row is not None else None
|
||||
|
||||
async def requeue_claimed_run(
|
||||
self,
|
||||
run_record_id: str,
|
||||
*,
|
||||
lease_owner: str,
|
||||
error: str | None = None,
|
||||
) -> bool:
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(
|
||||
update(ScheduledTaskRunRow)
|
||||
.where(
|
||||
ScheduledTaskRunRow.id == run_record_id,
|
||||
ScheduledTaskRunRow.status == "launching",
|
||||
ScheduledTaskRunRow.lease_owner == lease_owner,
|
||||
)
|
||||
.values(
|
||||
status="queued",
|
||||
lease_owner=None,
|
||||
lease_expires_at=None,
|
||||
error=error,
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
return result.rowcount == 1
|
||||
|
||||
async def expire_queued_runs(
|
||||
self,
|
||||
*,
|
||||
created_before: datetime,
|
||||
error: str,
|
||||
now: datetime,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Expire waiting rows and update each parent in one transaction.
|
||||
|
||||
The queued child keeps the task unclaimable until the parent lock is
|
||||
held. Releasing the child slot and advancing the parent in the same
|
||||
transaction prevents a peer scheduler from taking a stale due-task
|
||||
lease between those two writes.
|
||||
"""
|
||||
async with self._sf() as session:
|
||||
candidate_keys = list(
|
||||
(
|
||||
await session.execute(
|
||||
select(
|
||||
ScheduledTaskRunRow.id,
|
||||
ScheduledTaskRunRow.task_id,
|
||||
)
|
||||
.where(
|
||||
ScheduledTaskRunRow.status == "queued",
|
||||
ScheduledTaskRunRow.created_at <= created_before,
|
||||
)
|
||||
.order_by(
|
||||
ScheduledTaskRunRow.task_id.asc(),
|
||||
ScheduledTaskRunRow.id.asc(),
|
||||
)
|
||||
)
|
||||
).all()
|
||||
)
|
||||
|
||||
expired: list[dict[str, Any]] = []
|
||||
for row_id, task_id in candidate_keys:
|
||||
async with self._sf() as session:
|
||||
task = await self._lock_task(session, task_id)
|
||||
row = await session.get(ScheduledTaskRunRow, row_id, with_for_update=True)
|
||||
if row is None or row.status != "queued":
|
||||
await session.rollback()
|
||||
continue
|
||||
|
||||
row.status = "failed"
|
||||
row.error = error
|
||||
row.finished_at = now
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
|
||||
if task is not None:
|
||||
if task.status == "paused":
|
||||
# A concurrent/later pause owns the parent presentation.
|
||||
task.lease_owner = None
|
||||
task.lease_expires_at = None
|
||||
else:
|
||||
if row.trigger == "manual":
|
||||
next_at = task.next_run_at
|
||||
task_status = task.status or "enabled"
|
||||
else:
|
||||
next_at = compute_next_run_at(
|
||||
task.schedule_type,
|
||||
task.schedule_spec,
|
||||
task.timezone,
|
||||
now=now,
|
||||
)
|
||||
task_status = "failed" if task.schedule_type == "once" else "enabled"
|
||||
task.status = task_status
|
||||
task.next_run_at = next_at
|
||||
task.last_thread_id = row.thread_id
|
||||
task.last_error = error
|
||||
task.lease_owner = None
|
||||
task.lease_expires_at = None
|
||||
task.updated_at = datetime.now(UTC)
|
||||
|
||||
await session.commit()
|
||||
expired.append(self._row_to_dict(row))
|
||||
return expired
|
||||
|
||||
async def fail_launching_run(
|
||||
self,
|
||||
run_record_id: str,
|
||||
*,
|
||||
task_id: str,
|
||||
lease_owner: str,
|
||||
error: str,
|
||||
now: datetime,
|
||||
) -> bool:
|
||||
"""Fail a claimed launch and update its parent without a release gap."""
|
||||
async with self._sf() as session:
|
||||
task = await self._lock_task(session, task_id)
|
||||
row = await session.get(ScheduledTaskRunRow, run_record_id, with_for_update=True)
|
||||
if row is None or row.task_id != task_id or row.status != "launching" or row.lease_owner != lease_owner:
|
||||
await session.rollback()
|
||||
return False
|
||||
|
||||
row.status = "failed"
|
||||
row.error = error
|
||||
row.started_at = now
|
||||
row.finished_at = now
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
|
||||
if task is not None:
|
||||
if row.trigger == "manual":
|
||||
task_status = task.status or "enabled"
|
||||
next_at = task.next_run_at
|
||||
else:
|
||||
task_status = "failed" if task.schedule_type == "once" else "enabled"
|
||||
next_at = compute_next_run_at(
|
||||
task.schedule_type,
|
||||
task.schedule_spec,
|
||||
task.timezone,
|
||||
now=now,
|
||||
)
|
||||
task.status = task_status
|
||||
task.next_run_at = next_at
|
||||
task.last_run_at = now
|
||||
task.last_run_id = None
|
||||
task.last_thread_id = row.thread_id
|
||||
task.last_error = error
|
||||
task.lease_owner = None
|
||||
task.lease_expires_at = None
|
||||
task.updated_at = datetime.now(UTC)
|
||||
|
||||
await session.commit()
|
||||
return True
|
||||
|
||||
async def reconcile_launched_run(
|
||||
self,
|
||||
run_record_id: str,
|
||||
*,
|
||||
task_id: str,
|
||||
run_id: str,
|
||||
started_at: datetime,
|
||||
) -> bool:
|
||||
"""Associate a run that committed after its short claim was recovered.
|
||||
|
||||
The Gateway launch path is idempotent per scheduled occurrence, so a
|
||||
peer that already reclaimed this row resolves to the same durable run.
|
||||
Parent-first locking restores the active slot before later parent
|
||||
bookkeeping can make the task due again.
|
||||
"""
|
||||
async with self._sf() as session:
|
||||
await self._lock_task(session, task_id)
|
||||
row = await session.get(ScheduledTaskRunRow, run_record_id, with_for_update=True)
|
||||
if row is None or row.task_id != task_id:
|
||||
await session.rollback()
|
||||
return False
|
||||
if row.status in TERMINAL_RUN_STATUSES:
|
||||
if row.run_id is None:
|
||||
# Timeout/pause can terminalize a claim that recovery
|
||||
# briefly returned to ``queued`` while Gateway admission
|
||||
# was still completing. A returned durable run is the
|
||||
# stronger fact; completion-owned terminal rows already
|
||||
# carry this exact run_id and stay terminal below.
|
||||
row.status = "running"
|
||||
row.run_id = run_id
|
||||
row.error = None
|
||||
row.finished_at = None
|
||||
row.started_at = row.started_at or started_at
|
||||
elif row.run_id != run_id:
|
||||
await session.rollback()
|
||||
return False
|
||||
elif row.started_at is None:
|
||||
row.started_at = started_at
|
||||
else:
|
||||
row.status = "running"
|
||||
row.run_id = run_id
|
||||
row.error = None
|
||||
row.started_at = row.started_at or started_at
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
await session.commit()
|
||||
return True
|
||||
|
||||
async def recover_expired_launch_claims(self, *, error: str, now: datetime) -> int:
|
||||
"""Recover single-instance claims that outlived their short lease."""
|
||||
stmt = (
|
||||
select(
|
||||
ScheduledTaskRunRow.id,
|
||||
ScheduledTaskRunRow.task_id,
|
||||
)
|
||||
.where(
|
||||
ScheduledTaskRunRow.status == "launching",
|
||||
or_(
|
||||
ScheduledTaskRunRow.lease_expires_at.is_(None),
|
||||
ScheduledTaskRunRow.lease_expires_at < now,
|
||||
),
|
||||
)
|
||||
.order_by(
|
||||
ScheduledTaskRunRow.task_id.asc(),
|
||||
ScheduledTaskRunRow.id.asc(),
|
||||
)
|
||||
)
|
||||
async with self._sf() as session:
|
||||
row_keys = list((await session.execute(stmt)).all())
|
||||
recovered = 0
|
||||
for row_id, task_id in row_keys:
|
||||
task = await self._lock_task(session, task_id)
|
||||
row = await session.get(ScheduledTaskRunRow, row_id, with_for_update=True)
|
||||
if row is None or row.status != "launching":
|
||||
continue
|
||||
candidate = await self._find_underlying_run(session, row, task)
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
if candidate is None:
|
||||
row.status = "queued"
|
||||
else:
|
||||
self._associate_scheduled_run(row, candidate)
|
||||
self._associate_task_with_run(task, row, candidate)
|
||||
if candidate.status in {"pending", "running"}:
|
||||
row.status = "running"
|
||||
row.error = None
|
||||
elif candidate.status == "success":
|
||||
row.status = "success"
|
||||
row.error = None
|
||||
row.finished_at = now
|
||||
elif candidate.status in {"error", "timeout"}:
|
||||
row.status = "failed"
|
||||
row.error = candidate.error
|
||||
row.finished_at = now
|
||||
else:
|
||||
row.status = "interrupted"
|
||||
row.error = candidate.error or error
|
||||
row.finished_at = now
|
||||
recovered += 1
|
||||
await session.commit()
|
||||
return recovered
|
||||
|
||||
async def update_status(
|
||||
self,
|
||||
run_record_id: str,
|
||||
@ -131,29 +602,45 @@ class ScheduledTaskRunRepository:
|
||||
started_at: datetime | None = None,
|
||||
finished_at: datetime | None = None,
|
||||
protect_terminal: bool = False,
|
||||
) -> None:
|
||||
expected_lease_owner: str | None = None,
|
||||
) -> bool:
|
||||
async with self._sf() as session:
|
||||
row = await session.get(ScheduledTaskRunRow, run_record_id)
|
||||
if row is None:
|
||||
return
|
||||
return False
|
||||
if protect_terminal and row.status in TERMINAL_RUN_STATUSES:
|
||||
# The launch-path "running" write lost the race against the
|
||||
# completion hook; keep the terminal status/error and only
|
||||
# backfill bookkeeping the completion write could not know.
|
||||
# Completion clears the short launch lease, so allow that
|
||||
# backfill after an owner mismatch only when the terminal row
|
||||
# already identifies the exact same durable run. A stale
|
||||
# launcher for another run remains fenced.
|
||||
same_run = run_id is not None and row.run_id == run_id
|
||||
if expected_lease_owner is not None and row.lease_owner != expected_lease_owner and not same_run:
|
||||
await session.rollback()
|
||||
return False
|
||||
if row.run_id is None and run_id is not None:
|
||||
row.run_id = run_id
|
||||
if row.started_at is None and started_at is not None:
|
||||
row.started_at = started_at
|
||||
await session.commit()
|
||||
return
|
||||
return True
|
||||
if expected_lease_owner is not None and row.lease_owner != expected_lease_owner:
|
||||
await session.rollback()
|
||||
return False
|
||||
row.status = status
|
||||
row.run_id = run_id
|
||||
row.error = error
|
||||
if status != "launching":
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
if started_at is not None:
|
||||
row.started_at = started_at
|
||||
if finished_at is not None:
|
||||
row.finished_at = finished_at
|
||||
await session.commit()
|
||||
return True
|
||||
|
||||
async def has_active_runs(self, task_id: str) -> bool:
|
||||
stmt = (
|
||||
@ -169,23 +656,54 @@ class ScheduledTaskRunRepository:
|
||||
return result.scalars().first() is not None
|
||||
|
||||
async def mark_stale_active_runs(self, *, error: str) -> int:
|
||||
"""Fail-fast bookkeeping for runs orphaned by a process crash.
|
||||
"""Recover single-instance launch claims and fail orphaned live runs.
|
||||
|
||||
Agent runs execute in-process, so any ``queued``/``running`` row found
|
||||
at scheduler startup belongs to a run whose process is gone. Only valid
|
||||
under the MVP's single-scheduler-instance assumption.
|
||||
Waiting rows are durable queue entries and survive restart. A
|
||||
``launching`` row without a committed live run is safe to retry; a
|
||||
``running`` row belonged to the dead in-process runtime.
|
||||
"""
|
||||
stmt = select(ScheduledTaskRunRow).where(ScheduledTaskRunRow.status.in_(ACTIVE_RUN_STATUSES))
|
||||
stmt = (
|
||||
select(
|
||||
ScheduledTaskRunRow.id,
|
||||
ScheduledTaskRunRow.task_id,
|
||||
)
|
||||
.where(ScheduledTaskRunRow.status.in_(EXECUTING_RUN_STATUSES))
|
||||
.order_by(
|
||||
ScheduledTaskRunRow.task_id.asc(),
|
||||
ScheduledTaskRunRow.id.asc(),
|
||||
)
|
||||
)
|
||||
now = datetime.now(UTC)
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(stmt)
|
||||
rows = list(result.scalars())
|
||||
for row in rows:
|
||||
row.status = "interrupted"
|
||||
row.error = error
|
||||
row.finished_at = now
|
||||
row_keys = list((await session.execute(stmt)).all())
|
||||
recovered = 0
|
||||
for row_id, task_id in row_keys:
|
||||
task = await self._lock_task(session, task_id)
|
||||
row = await session.get(ScheduledTaskRunRow, row_id, with_for_update=True)
|
||||
if row is None or row.status not in EXECUTING_RUN_STATUSES:
|
||||
continue
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
candidate = await self._find_underlying_run(session, row, task)
|
||||
if row.status == "launching" and candidate is None:
|
||||
row.status = "queued"
|
||||
else:
|
||||
if candidate is not None:
|
||||
self._associate_scheduled_run(row, candidate)
|
||||
self._associate_task_with_run(task, row, candidate)
|
||||
if candidate is not None and candidate.status == "success":
|
||||
row.status = "success"
|
||||
row.error = None
|
||||
elif candidate is not None and candidate.status in {"error", "timeout"}:
|
||||
row.status = "failed"
|
||||
row.error = candidate.error
|
||||
else:
|
||||
row.status = "interrupted"
|
||||
row.error = error
|
||||
row.finished_at = now
|
||||
recovered += 1
|
||||
await session.commit()
|
||||
return len(rows)
|
||||
return recovered
|
||||
|
||||
async def reconcile_active_runs(
|
||||
self,
|
||||
@ -201,17 +719,67 @@ class ScheduledTaskRunRepository:
|
||||
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())
|
||||
result = await session.execute(
|
||||
select(
|
||||
ScheduledTaskRunRow.id,
|
||||
ScheduledTaskRunRow.task_id,
|
||||
)
|
||||
.where(ScheduledTaskRunRow.status.in_(EXECUTING_RUN_STATUSES))
|
||||
.order_by(
|
||||
ScheduledTaskRunRow.task_id.asc(),
|
||||
ScheduledTaskRunRow.id.asc(),
|
||||
)
|
||||
)
|
||||
row_keys = list(result.all())
|
||||
stale = 0
|
||||
for row_id in row_ids:
|
||||
associations: list[tuple[ScheduledTaskRow | None, ScheduledTaskRunRow, RunRow]] = []
|
||||
for row_id, task_id in row_keys:
|
||||
# Keep the same task -> scheduled-run lock order used by
|
||||
# pause/delete. Reversing these two locks lets a user action
|
||||
# and a peer reconciliation deadlock each other on Postgres.
|
||||
# Multi-instance reconciliation is Postgres-only. Keep this a
|
||||
# row lock without SQLite's writer-lock emulation because the
|
||||
# SQLite regression path performs durable-run takeover in a
|
||||
# nested short transaction below.
|
||||
task = await session.get(ScheduledTaskRow, task_id, with_for_update=True)
|
||||
row = await session.get(ScheduledTaskRunRow, row_id, with_for_update=True)
|
||||
if row is None or row.status not in ACTIVE_RUN_STATUSES:
|
||||
if row is None or row.status not in EXECUTING_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:
|
||||
self._associate_scheduled_run(row, candidate)
|
||||
# Defer parent writes until all run takeovers have finished.
|
||||
# Flushing a parent mutation before claim_for_takeover()
|
||||
# would hold SQLite's writer lock across the nested short
|
||||
# transaction used by that durable-run CAS.
|
||||
associations.append((task, row, candidate))
|
||||
if candidate is not None and candidate.status not in {"pending", "running"}:
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
if candidate.status == "success":
|
||||
row.status = "success"
|
||||
row.error = None
|
||||
elif candidate.status in {"error", "timeout"}:
|
||||
row.status = "failed"
|
||||
row.error = candidate.error
|
||||
else:
|
||||
row.status = "interrupted"
|
||||
row.error = candidate.error or error
|
||||
row.finished_at = now
|
||||
stale += 1
|
||||
continue
|
||||
if candidate is not None and candidate.status in {"pending", "running"}:
|
||||
if _lease_is_alive(candidate.lease_expires_at, now=now, grace_seconds=lease_grace_seconds):
|
||||
# A peer can observe the committed durable run before
|
||||
# the launcher writes its scheduled-run bookkeeping.
|
||||
# Once reconciliation releases the short launch claim,
|
||||
# that owner-fenced write must be allowed to fail
|
||||
# without leaving incomplete or stale history behind.
|
||||
row.error = None
|
||||
if row.status == "launching":
|
||||
row.status = "running"
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
continue
|
||||
# Run takeover commits in its own short transaction. If this
|
||||
# outer commit fails, the next poll finishes scheduled-row
|
||||
@ -226,12 +794,22 @@ class ScheduledTaskRunRepository:
|
||||
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):
|
||||
if row.status == "launching" and row.run_id is None:
|
||||
if _lease_is_alive(row.lease_expires_at, now=now, grace_seconds=0):
|
||||
continue
|
||||
row.status = "queued"
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
stale += 1
|
||||
continue
|
||||
row.status = "interrupted"
|
||||
row.error = error
|
||||
row.finished_at = now
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
stale += 1
|
||||
for task, row, candidate in associations:
|
||||
self._associate_task_with_run(task, row, candidate)
|
||||
await session.commit()
|
||||
return stale
|
||||
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
from .model import ScheduledTaskRow
|
||||
from .sql import ScheduledTaskRepository
|
||||
from .sql import ActiveScheduledTaskMutationConflict, ScheduledTaskRepository
|
||||
|
||||
__all__ = ["ScheduledTaskRow", "ScheduledTaskRepository"]
|
||||
__all__ = ["ActiveScheduledTaskMutationConflict", "ScheduledTaskRow", "ScheduledTaskRepository"]
|
||||
|
||||
@ -22,7 +22,7 @@ class ScheduledTaskRow(Base):
|
||||
schedule_spec: Mapped[dict] = mapped_column(JSON, default=dict)
|
||||
timezone: Mapped[str] = mapped_column(String(64))
|
||||
status: Mapped[str] = mapped_column(String(16), default="enabled", index=True)
|
||||
overlap_policy: Mapped[str] = mapped_column(String(16), default="skip")
|
||||
overlap_policy: Mapped[str] = mapped_column(String(16), default="enqueue")
|
||||
next_run_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), index=True, nullable=True)
|
||||
last_run_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
last_run_id: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
||||
|
||||
@ -4,7 +4,7 @@ import logging
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import and_, exists, func, or_, select, text
|
||||
from sqlalchemy import and_, exists, or_, select, update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
from deerflow.persistence.run import RunRepository
|
||||
@ -16,7 +16,14 @@ 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
|
||||
|
||||
|
||||
class ActiveScheduledTaskMutationConflict(Exception):
|
||||
"""A user mutation raced an admitted scheduled-task occurrence."""
|
||||
|
||||
def __init__(self, status: str) -> None:
|
||||
self.status = status
|
||||
super().__init__(f"scheduled task has an active {status} occurrence")
|
||||
|
||||
|
||||
def _lease_is_alive(lease_expires_at: datetime | None, *, now: datetime, grace_seconds: int = 0) -> bool:
|
||||
@ -69,6 +76,14 @@ class ScheduledTaskRepository:
|
||||
data[key] = coerce_iso(data[key])
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
async def _lock_task(session: AsyncSession, task_id: str) -> ScheduledTaskRow | None:
|
||||
# Match scheduled-run admission on SQLite, where FOR UPDATE is ignored:
|
||||
# acquire the database writer before checking the child occurrence.
|
||||
if session.get_bind().dialect.name == "sqlite":
|
||||
await session.execute(update(ScheduledTaskRow).where(ScheduledTaskRow.id == task_id).values(updated_at=ScheduledTaskRow.updated_at))
|
||||
return await session.get(ScheduledTaskRow, task_id, with_for_update=True)
|
||||
|
||||
async def create(
|
||||
self,
|
||||
*,
|
||||
@ -113,23 +128,144 @@ class ScheduledTaskRepository:
|
||||
return None
|
||||
return self._row_to_dict(row)
|
||||
|
||||
async def get_internal(self, task_id: str) -> dict[str, Any] | None:
|
||||
"""Load a task for the internal queue worker without an auth boundary."""
|
||||
async with self._sf() as session:
|
||||
row = await session.get(ScheduledTaskRow, task_id)
|
||||
return self._row_to_dict(row) if row is not None else None
|
||||
|
||||
async def list_by_user(self, user_id: str) -> list[dict[str, Any]]:
|
||||
stmt = select(ScheduledTaskRow).where(ScheduledTaskRow.user_id == user_id).order_by(ScheduledTaskRow.created_at.desc(), ScheduledTaskRow.id.desc())
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(stmt)
|
||||
return [self._row_to_dict(row) for row in result.scalars()]
|
||||
|
||||
async def get_active_run_status(self, task_id: str) -> str | None:
|
||||
stmt = (
|
||||
select(ScheduledTaskRunRow.status)
|
||||
.where(
|
||||
ScheduledTaskRunRow.task_id == task_id,
|
||||
ScheduledTaskRunRow.status.in_(("queued", "launching", "running")),
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
async with self._sf() as session:
|
||||
return (await session.execute(stmt)).scalars().first()
|
||||
|
||||
async def pause_with_queue_cancellation(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
user_id: str,
|
||||
error: str,
|
||||
now: datetime,
|
||||
) -> str:
|
||||
"""Pause a task and cancel its waiting occurrence in one transaction."""
|
||||
async with self._sf() as session:
|
||||
task = await self._lock_task(session, task_id)
|
||||
if task is None or task.user_id != user_id:
|
||||
await session.rollback()
|
||||
return "not_found"
|
||||
run = (
|
||||
(
|
||||
await session.execute(
|
||||
select(ScheduledTaskRunRow)
|
||||
.where(
|
||||
ScheduledTaskRunRow.task_id == task_id,
|
||||
ScheduledTaskRunRow.status.in_(("queued", "launching", "running")),
|
||||
)
|
||||
.limit(1)
|
||||
.with_for_update()
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.first()
|
||||
)
|
||||
if run is not None and run.status in {"launching", "running"}:
|
||||
await session.rollback()
|
||||
return "executing"
|
||||
if run is not None:
|
||||
run.status = "interrupted"
|
||||
run.error = error
|
||||
run.finished_at = now
|
||||
run.lease_owner = None
|
||||
run.lease_expires_at = None
|
||||
task.status = "paused"
|
||||
task.lease_owner = None
|
||||
task.lease_expires_at = None
|
||||
task.updated_at = now
|
||||
await session.commit()
|
||||
return "paused"
|
||||
|
||||
async def delete_with_queue_cancellation(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
user_id: str,
|
||||
error: str,
|
||||
now: datetime,
|
||||
) -> str:
|
||||
"""Delete a task only before queue execution begins."""
|
||||
async with self._sf() as session:
|
||||
task = await self._lock_task(session, task_id)
|
||||
if task is None or task.user_id != user_id:
|
||||
await session.rollback()
|
||||
return "not_found"
|
||||
run = (
|
||||
(
|
||||
await session.execute(
|
||||
select(ScheduledTaskRunRow)
|
||||
.where(
|
||||
ScheduledTaskRunRow.task_id == task_id,
|
||||
ScheduledTaskRunRow.status.in_(("queued", "launching", "running")),
|
||||
)
|
||||
.limit(1)
|
||||
.with_for_update()
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.first()
|
||||
)
|
||||
if run is not None and run.status in {"launching", "running"}:
|
||||
await session.rollback()
|
||||
return "executing"
|
||||
if run is not None:
|
||||
run.status = "interrupted"
|
||||
run.error = error
|
||||
run.finished_at = now
|
||||
run.lease_owner = None
|
||||
run.lease_expires_at = None
|
||||
await session.delete(task)
|
||||
await session.commit()
|
||||
return "deleted"
|
||||
|
||||
async def update(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
user_id: str,
|
||||
updates: dict[str, Any],
|
||||
require_mutable: bool = False,
|
||||
) -> dict[str, Any] | None:
|
||||
async with self._sf() as session:
|
||||
row = await session.get(ScheduledTaskRow, task_id)
|
||||
row = await self._lock_task(session, task_id) if require_mutable else await session.get(ScheduledTaskRow, task_id)
|
||||
if row is None or row.user_id != user_id:
|
||||
return None
|
||||
if require_mutable:
|
||||
if row.status == "running":
|
||||
await session.rollback()
|
||||
raise ActiveScheduledTaskMutationConflict("running")
|
||||
active_status = await session.scalar(
|
||||
select(ScheduledTaskRunRow.status)
|
||||
.where(
|
||||
ScheduledTaskRunRow.task_id == task_id,
|
||||
ScheduledTaskRunRow.status.in_(("queued", "launching", "running")),
|
||||
)
|
||||
.limit(1)
|
||||
)
|
||||
if active_status is not None:
|
||||
await session.rollback()
|
||||
raise ActiveScheduledTaskMutationConflict(active_status)
|
||||
for key, value in updates.items():
|
||||
if hasattr(row, key):
|
||||
setattr(row, key, value)
|
||||
@ -154,31 +290,15 @@ 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)
|
||||
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")),
|
||||
)
|
||||
active_run_for_task = exists(
|
||||
select(ScheduledTaskRunRow.id).where(
|
||||
ScheduledTaskRunRow.task_id == ScheduledTaskRow.id,
|
||||
ScheduledTaskRunRow.status.in_(("queued", "launching", "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 = (
|
||||
@ -186,6 +306,7 @@ class ScheduledTaskRepository:
|
||||
.where(
|
||||
ScheduledTaskRow.next_run_at.is_not(None),
|
||||
ScheduledTaskRow.next_run_at <= now,
|
||||
~active_run_for_task,
|
||||
or_(
|
||||
and_(
|
||||
ScheduledTaskRow.status == "enabled",
|
||||
@ -218,12 +339,57 @@ class ScheduledTaskRepository:
|
||||
await session.commit()
|
||||
return [self._row_to_dict(row) for row in rows]
|
||||
|
||||
async def release_dispatch_lease(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
expected_lease_owner: str | None,
|
||||
status: str,
|
||||
) -> bool:
|
||||
"""Release the short due-task claim after its occurrence is queued."""
|
||||
async with self._sf() as session:
|
||||
row = await session.get(ScheduledTaskRow, task_id, with_for_update=True)
|
||||
if row is None:
|
||||
return False
|
||||
if expected_lease_owner is not None and row.lease_owner != expected_lease_owner:
|
||||
await session.rollback()
|
||||
return False
|
||||
row.status = status
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
row.updated_at = datetime.now(UTC)
|
||||
await session.commit()
|
||||
return True
|
||||
|
||||
async def release_queued_admission_lease(self, task_id: str) -> bool:
|
||||
"""Recover a crash after queue insert but before parent-lease release."""
|
||||
async with self._sf() as session:
|
||||
task = await session.get(ScheduledTaskRow, task_id, with_for_update=True)
|
||||
if task is None or task.status != "running" or task.lease_owner is None:
|
||||
await session.rollback()
|
||||
return False
|
||||
queued = await session.scalar(
|
||||
select(ScheduledTaskRunRow.id).where(
|
||||
ScheduledTaskRunRow.task_id == task_id,
|
||||
ScheduledTaskRunRow.status == "queued",
|
||||
)
|
||||
)
|
||||
if queued is None:
|
||||
await session.rollback()
|
||||
return False
|
||||
task.status = "enabled"
|
||||
task.lease_owner = None
|
||||
task.lease_expires_at = None
|
||||
task.updated_at = datetime.now(UTC)
|
||||
await session.commit()
|
||||
return True
|
||||
|
||||
async def update_after_launch(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
status: str,
|
||||
next_run_at: datetime | None,
|
||||
next_run_at: datetime | str | None,
|
||||
last_run_at: datetime | str | None,
|
||||
last_run_id: str | None,
|
||||
last_thread_id: str | None,
|
||||
@ -254,11 +420,12 @@ class ScheduledTaskRepository:
|
||||
else:
|
||||
row.status = status
|
||||
row.last_error = last_error
|
||||
row.next_run_at = next_run_at
|
||||
should_increment_run_count = increment_run_count and (last_run_id is None or row.last_run_id != last_run_id)
|
||||
row.next_run_at = _coerce_datetime(next_run_at)
|
||||
row.last_run_at = _coerce_datetime(last_run_at)
|
||||
row.last_run_id = last_run_id
|
||||
row.last_thread_id = last_thread_id
|
||||
if increment_run_count:
|
||||
if should_increment_run_count:
|
||||
row.run_count += 1
|
||||
row.lease_owner = None
|
||||
row.lease_expires_at = None
|
||||
@ -363,7 +530,7 @@ class ScheduledTaskRepository:
|
||||
select(ScheduledTaskRunRow)
|
||||
.where(
|
||||
ScheduledTaskRunRow.task_id == task.id,
|
||||
ScheduledTaskRunRow.status.in_(("queued", "running")),
|
||||
ScheduledTaskRunRow.status.in_(("queued", "launching", "running")),
|
||||
)
|
||||
.order_by(ScheduledTaskRunRow.created_at.desc())
|
||||
.limit(1)
|
||||
|
||||
@ -2250,11 +2250,12 @@ def test_launch_scheduled_thread_run_marks_context_non_interactive(_stub_app_con
|
||||
async def _scenario():
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def fake_start_run(body, thread_id, request):
|
||||
async def fake_start_run(body, thread_id, request, *, idempotency_key=None):
|
||||
captured["body"] = body
|
||||
captured["thread_id"] = thread_id
|
||||
captured["context"] = body.context
|
||||
captured["metadata"] = body.metadata
|
||||
captured["idempotency_key"] = idempotency_key
|
||||
captured["if_not_exists"] = body.if_not_exists
|
||||
captured["on_completion"] = body.on_completion
|
||||
return SimpleNamespace(run_id="run-1", thread_id=thread_id)
|
||||
@ -2266,7 +2267,10 @@ def test_launch_scheduled_thread_run_marks_context_non_interactive(_stub_app_con
|
||||
prompt="Run in background",
|
||||
app=SimpleNamespace(state=SimpleNamespace()),
|
||||
owner_user_id="user-1",
|
||||
metadata={"scheduled_task_id": "task-1"},
|
||||
metadata={
|
||||
"scheduled_task_id": "task-1",
|
||||
"scheduled_task_run_id": "task-run-1",
|
||||
},
|
||||
)
|
||||
return captured, result
|
||||
|
||||
@ -2276,7 +2280,11 @@ def test_launch_scheduled_thread_run_marks_context_non_interactive(_stub_app_con
|
||||
assert isinstance(captured["body"], RunCreateRequest)
|
||||
assert captured["body"].config == {"recursion_limit": 1000}
|
||||
assert captured["context"] == {"non_interactive": True, "user_id": "user-1"}
|
||||
assert captured["metadata"] == {"scheduled_task_id": "task-1"}
|
||||
assert captured["metadata"] == {
|
||||
"scheduled_task_id": "task-1",
|
||||
"scheduled_task_run_id": "task-run-1",
|
||||
}
|
||||
assert captured["idempotency_key"] == "scheduled-task:task-run-1"
|
||||
assert captured["if_not_exists"] == "create"
|
||||
assert captured["on_completion"] is None
|
||||
assert result == {"run_id": "run-1", "thread_id": "thread-scheduled"}
|
||||
@ -2302,7 +2310,8 @@ def test_launch_scheduled_thread_run_uses_configured_recursion_limit(_stub_app_c
|
||||
async def _scenario():
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def fake_start_run(body, thread_id, request):
|
||||
async def fake_start_run(body, thread_id, request, *, idempotency_key=None):
|
||||
assert idempotency_key is None
|
||||
captured["config"] = body.config
|
||||
return SimpleNamespace(run_id="run-1", thread_id=thread_id)
|
||||
|
||||
@ -2343,7 +2352,8 @@ def test_launch_scheduled_thread_run_recursion_limit_is_clamped_to_ceiling(_stub
|
||||
async def _scenario():
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def fake_start_run(body, thread_id, request):
|
||||
async def fake_start_run(body, thread_id, request, *, idempotency_key=None):
|
||||
assert idempotency_key is None
|
||||
captured["config"] = body.config
|
||||
return SimpleNamespace(run_id="run-1", thread_id=thread_id)
|
||||
|
||||
@ -2375,7 +2385,8 @@ def test_launch_scheduled_thread_run_falls_back_when_config_unloadable(_stub_app
|
||||
async def _scenario():
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def fake_start_run(body, thread_id, request):
|
||||
async def fake_start_run(body, thread_id, request, *, idempotency_key=None):
|
||||
assert idempotency_key is None
|
||||
captured["config"] = body.config
|
||||
return SimpleNamespace(run_id="run-1", thread_id=thread_id)
|
||||
|
||||
|
||||
@ -157,7 +157,7 @@ async def test_migration_dedupes_duplicate_active_rows_before_unique_index(tmp_p
|
||||
with sqlite3.connect(db_path) as raw:
|
||||
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
||||
# Bootstrap upgrades through the later revisions after 0004.
|
||||
assert version_row[0] == "0014_managed_subagents"
|
||||
assert version_row[0] == "0015_scheduled_task_enqueue"
|
||||
|
||||
# Sanity: the invariant the index enforces is now true — at most one
|
||||
# active row per thread.
|
||||
|
||||
@ -13,9 +13,10 @@ End-to-end shape (mirrors ``test_migration_0004_run_ownership_dedupe``):
|
||||
4. Run ``init_engine`` (the FastAPI lifespan entry point), which routes through
|
||||
``bootstrap_schema`` -> ``upgrade head`` -> ``0007.upgrade()``.
|
||||
5. Verify the migration superseded the older duplicates (set them to
|
||||
``interrupted`` with an explanatory message + ``finished_at``), kept the
|
||||
newest active row, and successfully built the ``uq_scheduled_task_run_active``
|
||||
partial unique index.
|
||||
``interrupted`` with an explanatory message + ``finished_at``), successfully
|
||||
built the ``uq_scheduled_task_run_active`` partial unique index, and let the
|
||||
later scheduler-enqueue migration normalize the surviving legacy ``queued``
|
||||
row before durable queue semantics become active.
|
||||
|
||||
Pre-fix codepath would have raised ``UNIQUE constraint failed`` (SQLite) /
|
||||
``could not create unique index`` (Postgres) on step 5, aborting the alembic
|
||||
@ -147,9 +148,12 @@ async def test_migration_supersedes_duplicate_active_runs_before_unique_index(tm
|
||||
try:
|
||||
runs = _fetch_runs(db_path)
|
||||
|
||||
# Newest active row on the duplicated task survives unchanged.
|
||||
assert runs["run-newest"][0] == "queued"
|
||||
assert runs["run-newest"][1] is None
|
||||
# 0007 keeps the newest duplicate, then 0015 deliberately interrupts
|
||||
# that pre-existing queued row before durable queue semantics begin.
|
||||
newest_status, newest_error, newest_finished_at = runs["run-newest"]
|
||||
assert newest_status == "interrupted"
|
||||
assert "gateway upgraded" in (newest_error or "")
|
||||
assert newest_finished_at is not None
|
||||
|
||||
# Older duplicate active rows are superseded with an explanatory error
|
||||
# and a finished_at timestamp (mark_stale_active_runs orphan semantics).
|
||||
@ -169,7 +173,7 @@ async def test_migration_supersedes_duplicate_active_runs_before_unique_index(tm
|
||||
|
||||
with sqlite3.connect(db_path) as raw:
|
||||
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
||||
assert version_row[0] == "0014_managed_subagents"
|
||||
assert version_row[0] == "0015_scheduled_task_enqueue"
|
||||
|
||||
# Sanity: the invariant the index enforces now holds — at most one
|
||||
# active row per task_id.
|
||||
|
||||
67
backend/tests/test_migration_0015_scheduled_task_enqueue.py
Normal file
67
backend/tests/test_migration_0015_scheduled_task_enqueue.py
Normal file
@ -0,0 +1,67 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
import deerflow.persistence.models # noqa: F401
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.bootstrap import bootstrap_schema
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
async def test_migration_interrupts_legacy_queue_and_adds_claim_fields(tmp_path: Path) -> None:
|
||||
db_path = tmp_path / "deer.db"
|
||||
sync = sa.create_engine(f"sqlite:///{db_path}")
|
||||
try:
|
||||
Base.metadata.create_all(sync)
|
||||
with sync.begin() as conn:
|
||||
conn.execute(sa.text("DROP INDEX uq_scheduled_task_run_active"))
|
||||
conn.execute(sa.text("ALTER TABLE scheduled_task_runs DROP COLUMN attempt_count"))
|
||||
conn.execute(sa.text("ALTER TABLE scheduled_task_runs DROP COLUMN lease_expires_at"))
|
||||
conn.execute(sa.text("ALTER TABLE scheduled_task_runs DROP COLUMN lease_owner"))
|
||||
conn.execute(sa.text("CREATE UNIQUE INDEX uq_scheduled_task_run_active ON scheduled_task_runs (task_id) WHERE status IN ('queued', 'running')"))
|
||||
conn.execute(
|
||||
sa.text(
|
||||
"INSERT INTO scheduled_tasks "
|
||||
"(id, user_id, thread_id, context_mode, assistant_id, title, prompt, "
|
||||
"schedule_type, schedule_spec, timezone, status, overlap_policy, run_count, created_at, updated_at) "
|
||||
"VALUES ('task-legacy', 'user-1', 'thread-1', 'reuse_thread', 'lead_agent', "
|
||||
"'Legacy', 'Prompt', 'cron', '{\"cron\":\"0 9 * * *\"}', 'UTC', "
|
||||
"'enabled', 'skip', 0, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
|
||||
)
|
||||
)
|
||||
conn.execute(
|
||||
sa.text(
|
||||
"INSERT INTO scheduled_task_runs (id, task_id, thread_id, scheduled_for, trigger, status, created_at) VALUES ('run-legacy-queued', 'task-legacy', 'thread-1', CURRENT_TIMESTAMP, 'scheduled', 'queued', CURRENT_TIMESTAMP)"
|
||||
)
|
||||
)
|
||||
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 ('0013_mcp_task_notifications')"))
|
||||
finally:
|
||||
sync.dispose()
|
||||
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}")
|
||||
try:
|
||||
await bootstrap_schema(engine, backend="sqlite")
|
||||
async with engine.connect() as conn:
|
||||
columns = {column["name"]: column for column in await conn.run_sync(lambda connection: sa.inspect(connection).get_columns("scheduled_task_runs"))}
|
||||
version = await conn.scalar(sa.text("SELECT version_num FROM alembic_version"))
|
||||
overlap_policy = await conn.scalar(sa.text("SELECT overlap_policy FROM scheduled_tasks WHERE id = 'task-legacy'"))
|
||||
legacy_run = (await conn.execute(sa.text("SELECT status, error, finished_at FROM scheduled_task_runs WHERE id = 'run-legacy-queued'"))).one()
|
||||
index_sql = await conn.scalar(sa.text("SELECT sql FROM sqlite_master WHERE type = 'index' AND name = 'uq_scheduled_task_run_active'"))
|
||||
|
||||
assert version == "0015_scheduled_task_enqueue"
|
||||
assert {"lease_owner", "lease_expires_at", "attempt_count"} <= columns.keys()
|
||||
assert columns["attempt_count"]["nullable"] is False
|
||||
assert overlap_policy == "enqueue"
|
||||
assert legacy_run.status == "interrupted"
|
||||
assert "gateway upgraded" in legacy_run.error
|
||||
assert legacy_run.finished_at is not None
|
||||
assert "'launching'" in index_sql
|
||||
finally:
|
||||
await engine.dispose()
|
||||
@ -48,7 +48,7 @@ from deerflow.persistence.migrations._helpers import _normalize_default
|
||||
asyncio_test = pytest.mark.asyncio
|
||||
|
||||
|
||||
HEAD = "0014_managed_subagents"
|
||||
HEAD = "0015_scheduled_task_enqueue"
|
||||
BASELINE = "0001_baseline"
|
||||
|
||||
|
||||
|
||||
@ -28,7 +28,7 @@ from deerflow.persistence.bootstrap import bootstrap_schema
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
HEAD = "0014_managed_subagents"
|
||||
HEAD = "0015_scheduled_task_enqueue"
|
||||
|
||||
|
||||
def _url(tmp_path: Path) -> str:
|
||||
|
||||
@ -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] == "0014_managed_subagents"
|
||||
assert version_row[0] == "0015_scheduled_task_enqueue"
|
||||
|
||||
# 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] == "0014_managed_subagents"
|
||||
assert version_row[0] == "0015_scheduled_task_enqueue"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
@ -1,15 +1,14 @@
|
||||
"""Concurrency regression tests for the scheduled-task dispatch TOCTOU.
|
||||
|
||||
``ScheduledTaskService.dispatch_task`` guards the "at most one active run per
|
||||
task when overlap_policy=skip" invariant with a non-atomic
|
||||
``has_active_runs`` check followed by a separate ``create(status="queued")``
|
||||
``ScheduledTaskService.dispatch_task`` guards the "at most one non-terminal
|
||||
occurrence per task" invariant with a non-atomic active-row lookup followed by
|
||||
a separate ``create(status="queued")``
|
||||
insert. Two concurrent dispatches (double-click, client retry, or a manual
|
||||
trigger racing the poller) can both pass the check and both launch. The fix
|
||||
makes the database the atomic arbiter via the partial unique index
|
||||
``uq_scheduled_task_run_active`` (``task_id WHERE status IN
|
||||
('queued','running')``); the losing insert is translated to the typed
|
||||
``ActiveScheduledRunConflict`` and collapsed to the same outcome as the
|
||||
fast-path check.
|
||||
('queued','launching','running')``); the losing insert is translated to the
|
||||
typed ``ActiveScheduledRunConflict``.
|
||||
|
||||
These tests drive the REAL ``ScheduledTaskRunRepository`` + ``ScheduledTaskService``
|
||||
against a real file-backed ``sqlite+aiosqlite`` DB (so the index is actually
|
||||
@ -31,22 +30,24 @@ from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
_ACTIVE_STATUSES = {"queued", "running"}
|
||||
_ACTIVE_STATUSES = {"queued", "launching", "running"}
|
||||
|
||||
|
||||
class _BarrierRunRepo(ScheduledTaskRunRepository):
|
||||
"""Real repository that only releases both dispatchers past
|
||||
``has_active_runs`` once both have read it, so their ``create()`` calls
|
||||
``get_active_run`` once both have read it, so their ``create()`` calls
|
||||
genuinely race for the task's single active slot — a deterministic
|
||||
reproduction of the check-then-insert TOCTOU."""
|
||||
|
||||
def __init__(self, session_factory, barrier: asyncio.Barrier | None) -> None:
|
||||
super().__init__(session_factory)
|
||||
self._barrier = barrier
|
||||
self._barrier_reads = 0
|
||||
|
||||
async def has_active_runs(self, task_id: str) -> bool:
|
||||
result = await super().has_active_runs(task_id)
|
||||
if self._barrier is not None:
|
||||
async def get_active_run(self, task_id: str):
|
||||
result = await super().get_active_run(task_id)
|
||||
self._barrier_reads += 1
|
||||
if self._barrier is not None and self._barrier_reads <= 2:
|
||||
await self._barrier.wait()
|
||||
return result
|
||||
|
||||
@ -87,7 +88,7 @@ async def _seed_task(task_repo: ScheduledTaskRepository, task_id: str) -> dict:
|
||||
)
|
||||
task = await task_repo.get(task_id, user_id="user-1")
|
||||
assert task is not None
|
||||
assert task["overlap_policy"] == "skip"
|
||||
assert task["overlap_policy"] == "enqueue"
|
||||
return task
|
||||
|
||||
|
||||
@ -114,13 +115,13 @@ async def test_two_concurrent_manual_dispatches_launch_exactly_once(tmp_path):
|
||||
)
|
||||
|
||||
outcomes = sorted(result["outcome"] for result in results)
|
||||
# Exactly one wins the active slot; the loser is a 409-style conflict.
|
||||
assert outcomes == ["conflict", "launched"], outcomes
|
||||
# Exactly one wins the occurrence slot; the loser either observes the
|
||||
# queued row and coalesces into it or sees execution already starting.
|
||||
assert outcomes.count("launched") == 1, outcomes
|
||||
assert set(outcomes) <= {"launched", "queued", "conflict"}, outcomes
|
||||
assert len(launched) == 1, launched
|
||||
assert await _active_run_count(run_repo, "task-race-manual") == 1
|
||||
# The manual loser records no run-history row (nothing was scheduled).
|
||||
conflict = next(r for r in results if r["outcome"] == "conflict")
|
||||
assert conflict["task_run_id"] is None
|
||||
assert len({r["task_run_id"] for r in results if r["task_run_id"] is not None}) == 1
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
@ -143,10 +144,9 @@ async def test_scheduled_and_manual_dispatch_launch_exactly_once(tmp_path):
|
||||
)
|
||||
|
||||
outcomes = sorted(result["outcome"] for result in results)
|
||||
# Whichever won launched; the loser is conflict (manual) or skipped
|
||||
# (scheduled). Which one wins is timing-dependent, but exactly one runs.
|
||||
# Whichever won launched; the loser coalesces or sees execution begin.
|
||||
assert outcomes.count("launched") == 1, outcomes
|
||||
assert set(outcomes) <= {"launched", "conflict", "skipped"}, outcomes
|
||||
assert set(outcomes) <= {"launched", "queued", "conflict"}, outcomes
|
||||
assert len(launched) == 1, launched
|
||||
assert await _active_run_count(run_repo, "task-race-mixed") == 1
|
||||
finally:
|
||||
|
||||
@ -13,6 +13,7 @@ def test_app_config_exposes_scheduler_section():
|
||||
assert config.scheduler.multi_instance is False
|
||||
assert config.scheduler.poll_interval_seconds == 5
|
||||
assert config.scheduler.lease_seconds == 120
|
||||
assert config.scheduler.queue_timeout_seconds == 3600
|
||||
|
||||
|
||||
def test_scheduled_task_models_registered():
|
||||
|
||||
@ -73,29 +73,39 @@ async def _create_cron_task(task_repo: ScheduledTaskRepository, task_id: str, *,
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_postgres_global_budget_serializes_cross_pod_claims(postgres_repositories):
|
||||
task_repo, _task_run_repo, _run_repo = 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)
|
||||
for suffix in ("a", "b"):
|
||||
await _create_cron_task(task_repo, f"task-{suffix}", next_run_at=now)
|
||||
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="queued",
|
||||
)
|
||||
|
||||
claims = await asyncio.gather(
|
||||
task_repo.claim_due_tasks(
|
||||
task_run_repo.claim_queued_run(
|
||||
"task-run-a",
|
||||
now=now,
|
||||
lease_owner="pod-a",
|
||||
lease_seconds=60,
|
||||
limit=1,
|
||||
global_max_concurrent_runs=1,
|
||||
),
|
||||
task_repo.claim_due_tasks(
|
||||
task_run_repo.claim_queued_run(
|
||||
"task-run-b",
|
||||
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
|
||||
assert sum(claimed is not None for claimed in claims) == 1
|
||||
statuses = {(await task_run_repo.list_by_task(task_id))[0]["status"] for task_id in ("task-a", "task-b")}
|
||||
assert statuses == {"queued", "launching"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
897
backend/tests/test_scheduled_task_queue.py
Normal file
897
backend/tests/test_scheduled_task_queue.py
Normal file
@ -0,0 +1,897 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import pytest
|
||||
|
||||
from app.scheduler.service import ScheduledTaskService
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
|
||||
from deerflow.persistence.run import RunRepository
|
||||
from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository
|
||||
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
|
||||
from deerflow.runtime import ConflictError
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
async def _seed_reuse_task(task_repo: ScheduledTaskRepository, *, task_id: str, now: datetime) -> dict:
|
||||
await task_repo.create(
|
||||
task_id=task_id,
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="Follow up",
|
||||
prompt="Continue from the existing conversation",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "0 9 * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=now + timedelta(days=1),
|
||||
)
|
||||
task = await task_repo.get(task_id, user_id="user-1")
|
||||
assert task is not None
|
||||
return task
|
||||
|
||||
|
||||
def _make_service(task_repo, run_repo, launch_run, *, queue_timeout_seconds: int = 3600):
|
||||
return ScheduledTaskService(
|
||||
task_repo=task_repo,
|
||||
task_run_repo=run_repo,
|
||||
launch_run=launch_run,
|
||||
poll_interval_seconds=5,
|
||||
lease_seconds=120,
|
||||
max_concurrent_runs=3,
|
||||
queue_timeout_seconds=queue_timeout_seconds,
|
||||
)
|
||||
|
||||
|
||||
async def test_busy_reuse_thread_is_queued_then_launched_on_a_later_poll(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
task_repo = ScheduledTaskRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
task = await _seed_reuse_task(task_repo, task_id="task-queue", now=now)
|
||||
attempts = 0
|
||||
|
||||
async def launch_run(**kwargs):
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
if attempts == 1:
|
||||
raise ConflictError("Thread thread-1 already has an active run")
|
||||
return {"run_id": "run-queued", "thread_id": kwargs["thread_id"]}
|
||||
|
||||
service = _make_service(task_repo, run_repo, launch_run)
|
||||
|
||||
result = await service.dispatch_task(task, now=now, trigger="manual")
|
||||
|
||||
assert result["outcome"] == "queued"
|
||||
rows = await run_repo.list_by_task("task-queue")
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["status"] == "queued"
|
||||
assert rows[0]["attempt_count"] == 1
|
||||
|
||||
await service.run_once(now=now + timedelta(seconds=5))
|
||||
|
||||
rows = await run_repo.list_by_task("task-queue")
|
||||
assert rows[0]["status"] == "running"
|
||||
assert rows[0]["run_id"] == "run-queued"
|
||||
assert rows[0]["attempt_count"] == 2
|
||||
assert attempts == 2
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_paused_task_manual_run_waits_for_busy_thread_and_stays_paused(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
task_repo = ScheduledTaskRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
await _seed_reuse_task(task_repo, task_id="task-paused-manual", now=now)
|
||||
await task_repo.update(
|
||||
"task-paused-manual",
|
||||
user_id="user-1",
|
||||
updates={"status": "paused"},
|
||||
)
|
||||
task = await task_repo.get("task-paused-manual", user_id="user-1")
|
||||
assert task is not None
|
||||
attempts = 0
|
||||
|
||||
async def launch_run(**kwargs):
|
||||
nonlocal attempts
|
||||
attempts += 1
|
||||
if attempts == 1:
|
||||
raise ConflictError("Thread thread-1 already has an active run")
|
||||
return {"run_id": "run-paused-manual", "thread_id": kwargs["thread_id"]}
|
||||
|
||||
service = _make_service(task_repo, run_repo, launch_run)
|
||||
|
||||
result = await service.dispatch_task(task, now=now, trigger="manual")
|
||||
|
||||
assert result["outcome"] == "queued"
|
||||
assert (await run_repo.list_by_task(task["id"]))[0]["status"] == "queued"
|
||||
assert (await task_repo.get(task["id"], user_id="user-1"))["status"] == "paused"
|
||||
|
||||
await service.run_once(now=now + timedelta(seconds=5))
|
||||
|
||||
row = (await run_repo.list_by_task(task["id"]))[0]
|
||||
assert row["status"] == "running"
|
||||
assert row["run_id"] == "run-paused-manual"
|
||||
assert attempts == 2
|
||||
assert (await task_repo.get(task["id"], user_id="user-1"))["status"] == "paused"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_queued_run_survives_single_instance_restart_sweep(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
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
now = datetime(2026, 8, 20, 9, 0, tzinfo=UTC)
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-queued",
|
||||
task_id="task-queued",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=now,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
|
||||
swept = await run_repo.mark_stale_active_runs(error="gateway restarted")
|
||||
|
||||
assert swept == 0
|
||||
assert (await run_repo.list_by_task("task-queued"))[0]["status"] == "queued"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_only_one_worker_can_claim_a_queued_run(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
run_repo_a = ScheduledTaskRunRepository(sf)
|
||||
run_repo_b = ScheduledTaskRunRepository(sf)
|
||||
now = datetime(2026, 8, 20, 9, 0, tzinfo=UTC)
|
||||
await run_repo_a.create(
|
||||
run_record_id="task-run-claim",
|
||||
task_id="task-claim",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=now,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
|
||||
claims = await asyncio.gather(
|
||||
run_repo_a.claim_queued_run(
|
||||
"task-run-claim",
|
||||
lease_owner="worker-a",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
),
|
||||
run_repo_b.claim_queued_run(
|
||||
"task-run-claim",
|
||||
lease_owner="worker-b",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
),
|
||||
)
|
||||
|
||||
assert sum(claim is not None for claim in claims) == 1
|
||||
row = (await run_repo_a.list_by_task("task-claim"))[0]
|
||||
assert row["status"] == "launching"
|
||||
assert row["attempt_count"] == 1
|
||||
assert row["lease_owner"] in {"worker-a", "worker-b"}
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_same_thread_queue_is_claimed_in_fifo_order(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
|
||||
run_repo_a = ScheduledTaskRunRepository(sf)
|
||||
run_repo_b = ScheduledTaskRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
for run_id, task_id in (
|
||||
("task-run-a", "task-a"),
|
||||
("task-run-b", "task-b"),
|
||||
):
|
||||
await run_repo_a.create(
|
||||
run_record_id=run_id,
|
||||
task_id=task_id,
|
||||
thread_id="shared-thread",
|
||||
scheduled_for=now,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
|
||||
newer = await run_repo_b.claim_queued_run(
|
||||
"task-run-b",
|
||||
lease_owner="worker-b",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
older = await run_repo_a.claim_queued_run(
|
||||
"task-run-a",
|
||||
lease_owner="worker-a",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
|
||||
assert newer is None
|
||||
assert older is not None
|
||||
assert await run_repo_b.list_queued_runs(limit=10) == []
|
||||
assert (
|
||||
await run_repo_b.claim_queued_run(
|
||||
"task-run-b",
|
||||
lease_owner="worker-b",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
is None
|
||||
)
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_queue_drain_rotates_busy_thread_heads_without_breaking_fifo(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
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
for run_id, task_id, thread_id in (
|
||||
("task-run-busy-old", "task-busy-old", "busy-thread"),
|
||||
("task-run-busy-new", "task-busy-new", "busy-thread"),
|
||||
("task-run-ready", "task-ready", "ready-thread"),
|
||||
):
|
||||
await run_repo.create(
|
||||
run_record_id=run_id,
|
||||
task_id=task_id,
|
||||
thread_id=thread_id,
|
||||
scheduled_for=now,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
|
||||
claimed = await run_repo.claim_queued_run(
|
||||
"task-run-busy-old",
|
||||
lease_owner="worker-a",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
assert claimed is not None
|
||||
assert await run_repo.requeue_claimed_run(
|
||||
"task-run-busy-old",
|
||||
lease_owner="worker-a",
|
||||
error="thread is busy",
|
||||
)
|
||||
|
||||
candidates = await run_repo.list_queued_runs(limit=10)
|
||||
|
||||
# The untouched ready thread gets the next bounded-drain slot, and the
|
||||
# newer row for the busy thread stays hidden behind its FIFO head.
|
||||
assert [row["id"] for row in candidates] == [
|
||||
"task-run-ready",
|
||||
"task-run-busy-old",
|
||||
]
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_expired_launch_claim_is_requeued_but_waiting_timeout_fails(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
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-stale-claim",
|
||||
task_id="task-stale-claim",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=now - timedelta(minutes=5),
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
claimed = await run_repo.claim_queued_run(
|
||||
"task-run-stale-claim",
|
||||
lease_owner="dead-worker",
|
||||
now=now,
|
||||
lease_seconds=5,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
assert claimed is not None
|
||||
|
||||
reconciled = await run_repo.reconcile_active_runs(error="gateway lease expired", now=now + timedelta(seconds=6))
|
||||
|
||||
assert reconciled == 1
|
||||
row = (await run_repo.list_by_task("task-stale-claim"))[0]
|
||||
assert row["status"] == "queued"
|
||||
assert row["lease_owner"] is None
|
||||
|
||||
expired = await run_repo.expire_queued_runs(
|
||||
created_before=now + timedelta(seconds=7),
|
||||
error="queue wait timeout exceeded",
|
||||
now=now + timedelta(seconds=7),
|
||||
)
|
||||
assert [item["id"] for item in expired] == ["task-run-stale-claim"]
|
||||
assert (await run_repo.list_by_task("task-stale-claim"))[0]["status"] == "failed"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_slow_launch_is_reassociated_after_lease_recovery(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
task_repo = ScheduledTaskRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
task = await _seed_reuse_task(task_repo, task_id="task-slow-launch", now=now)
|
||||
queued = await run_repo.create(
|
||||
run_record_id="task-run-slow-launch",
|
||||
task_id=task["id"],
|
||||
thread_id="thread-1",
|
||||
scheduled_for=now,
|
||||
trigger="manual",
|
||||
status="queued",
|
||||
)
|
||||
launch_started = asyncio.Event()
|
||||
allow_launch_return = asyncio.Event()
|
||||
|
||||
async def launch_run(**kwargs):
|
||||
launch_started.set()
|
||||
await allow_launch_return.wait()
|
||||
return {"run_id": "run-slow-launch", "thread_id": kwargs["thread_id"]}
|
||||
|
||||
service = ScheduledTaskService(
|
||||
task_repo=task_repo,
|
||||
task_run_repo=run_repo,
|
||||
launch_run=launch_run,
|
||||
poll_interval_seconds=1,
|
||||
lease_seconds=5,
|
||||
max_concurrent_runs=3,
|
||||
)
|
||||
attempt = asyncio.create_task(service._attempt_queued_run(task, queued, now=now))
|
||||
await launch_started.wait()
|
||||
|
||||
recovered = await run_repo.recover_expired_launch_claims(
|
||||
error="launch lease expired",
|
||||
now=now + timedelta(seconds=6),
|
||||
)
|
||||
assert recovered == 1
|
||||
assert (await run_repo.list_by_task(task["id"]))[0]["status"] == "queued"
|
||||
|
||||
allow_launch_return.set()
|
||||
result = await attempt
|
||||
|
||||
row = (await run_repo.list_by_task(task["id"]))[0]
|
||||
assert result["outcome"] == "launched"
|
||||
assert row["status"] == "running"
|
||||
assert row["run_id"] == "run-slow-launch"
|
||||
assert row["lease_owner"] is None
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_failed_launch_releases_child_and_advances_parent_atomically(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
task_repo = ScheduledTaskRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
await task_repo.create(
|
||||
task_id="task-failed-launch-atomic",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="Atomic failure",
|
||||
prompt="fail",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "0 9 * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=now - timedelta(minutes=1),
|
||||
)
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-failed-launch-atomic",
|
||||
task_id="task-failed-launch-atomic",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=now,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
assert (
|
||||
await run_repo.claim_queued_run(
|
||||
"task-run-failed-launch-atomic",
|
||||
lease_owner="worker-a",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
is not None
|
||||
)
|
||||
|
||||
failed, peer_claims = await asyncio.gather(
|
||||
run_repo.fail_launching_run(
|
||||
"task-run-failed-launch-atomic",
|
||||
task_id="task-failed-launch-atomic",
|
||||
lease_owner="worker-a",
|
||||
error="launch failed",
|
||||
now=now,
|
||||
),
|
||||
task_repo.claim_due_tasks(
|
||||
now=now,
|
||||
lease_owner="worker-b",
|
||||
lease_seconds=120,
|
||||
limit=1,
|
||||
),
|
||||
)
|
||||
|
||||
assert failed is True
|
||||
assert peer_claims == []
|
||||
row = (await run_repo.list_by_task("task-failed-launch-atomic"))[0]
|
||||
task = await task_repo.get("task-failed-launch-atomic", user_id="user-1")
|
||||
assert row["status"] == "failed"
|
||||
assert task is not None
|
||||
assert datetime.fromisoformat(task["next_run_at"]) > now
|
||||
assert task["lease_owner"] is None
|
||||
assert task["last_error"] == "launch failed"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_failed_manual_launch_preserves_paused_next_run_at(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
task_repo = ScheduledTaskRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
next_run_at = now + timedelta(days=1)
|
||||
await task_repo.create(
|
||||
task_id="task-failed-manual-launch",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="Manual launch failure",
|
||||
prompt="fail",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "0 9 * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=next_run_at,
|
||||
)
|
||||
await task_repo.update(
|
||||
"task-failed-manual-launch",
|
||||
user_id="user-1",
|
||||
updates={"status": "paused"},
|
||||
)
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-failed-manual-launch",
|
||||
task_id="task-failed-manual-launch",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=now,
|
||||
trigger="manual",
|
||||
status="queued",
|
||||
)
|
||||
assert (
|
||||
await run_repo.claim_queued_run(
|
||||
"task-run-failed-manual-launch",
|
||||
lease_owner="worker-a",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
is not None
|
||||
)
|
||||
|
||||
assert await run_repo.fail_launching_run(
|
||||
"task-run-failed-manual-launch",
|
||||
task_id="task-failed-manual-launch",
|
||||
lease_owner="worker-a",
|
||||
error="launch failed",
|
||||
now=now,
|
||||
)
|
||||
|
||||
row = (await run_repo.list_by_task("task-failed-manual-launch"))[0]
|
||||
task = await task_repo.get("task-failed-manual-launch", user_id="user-1")
|
||||
assert row["status"] == "failed"
|
||||
assert task is not None
|
||||
assert task["status"] == "paused"
|
||||
assert task["next_run_at"] == next_run_at.isoformat()
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_pause_atomically_cancels_waiting_run_but_rejects_launching_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)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
await _seed_reuse_task(task_repo, task_id="task-pause-queued", now=now)
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-pause-queued",
|
||||
task_id="task-pause-queued",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=now,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
|
||||
result = await task_repo.pause_with_queue_cancellation(
|
||||
"task-pause-queued",
|
||||
user_id="user-1",
|
||||
error="paused while queued",
|
||||
now=now,
|
||||
)
|
||||
|
||||
assert result == "paused"
|
||||
assert (await task_repo.get("task-pause-queued", user_id="user-1"))["status"] == "paused"
|
||||
assert (await run_repo.list_by_task("task-pause-queued"))[0]["status"] == "interrupted"
|
||||
|
||||
await _seed_reuse_task(task_repo, task_id="task-pause-launching", now=now)
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-pause-launching",
|
||||
task_id="task-pause-launching",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=now,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
await run_repo.claim_queued_run(
|
||||
"task-run-pause-launching",
|
||||
lease_owner="worker-a",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
|
||||
result = await task_repo.pause_with_queue_cancellation(
|
||||
"task-pause-launching",
|
||||
user_id="user-1",
|
||||
error="paused while queued",
|
||||
now=now,
|
||||
)
|
||||
|
||||
assert result == "executing"
|
||||
assert (await task_repo.get("task-pause-launching", user_id="user-1"))["status"] == "enabled"
|
||||
assert (await run_repo.list_by_task("task-pause-launching"))[0]["status"] == "launching"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_manual_enqueue_cannot_unpause_a_task_that_cancels_its_waiting_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)
|
||||
|
||||
class DelayedCreateRunRepository(ScheduledTaskRunRepository):
|
||||
def __init__(self, session_factory):
|
||||
super().__init__(session_factory)
|
||||
self.created = asyncio.Event()
|
||||
self.resume_dispatch = asyncio.Event()
|
||||
|
||||
async def create(self, **kwargs):
|
||||
row = await super().create(**kwargs)
|
||||
self.created.set()
|
||||
await self.resume_dispatch.wait()
|
||||
return row
|
||||
|
||||
run_repo = DelayedCreateRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
task = await _seed_reuse_task(task_repo, task_id="task-manual-pause-race", now=now)
|
||||
|
||||
async def launch_run(**_kwargs):
|
||||
raise AssertionError("a queued occurrence cancelled by pause must not launch")
|
||||
|
||||
service = _make_service(task_repo, run_repo, launch_run)
|
||||
dispatch = asyncio.create_task(service.dispatch_task(task, now=now, trigger="manual"))
|
||||
await run_repo.created.wait()
|
||||
|
||||
pause_result = await task_repo.pause_with_queue_cancellation(
|
||||
task["id"],
|
||||
user_id="user-1",
|
||||
error="paused while queued",
|
||||
now=now,
|
||||
)
|
||||
run_repo.resume_dispatch.set()
|
||||
result = await dispatch
|
||||
|
||||
assert pause_result == "paused"
|
||||
assert result["outcome"] == "queued"
|
||||
assert (await task_repo.get(task["id"], user_id="user-1"))["status"] == "paused"
|
||||
assert (await run_repo.list_by_task(task["id"]))[0]["status"] == "interrupted"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("action", ["pause", "delete"])
|
||||
async def test_pause_or_delete_before_manual_admission_prevents_launch(tmp_path, action):
|
||||
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)
|
||||
|
||||
class DelayedAdmissionRunRepository(ScheduledTaskRunRepository):
|
||||
def __init__(self, session_factory):
|
||||
super().__init__(session_factory)
|
||||
self.before_insert = asyncio.Event()
|
||||
self.resume_insert = asyncio.Event()
|
||||
|
||||
async def create(self, **kwargs):
|
||||
self.before_insert.set()
|
||||
await self.resume_insert.wait()
|
||||
return await super().create(**kwargs)
|
||||
|
||||
run_repo = DelayedAdmissionRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
task = await _seed_reuse_task(task_repo, task_id=f"task-{action}-before-admission", now=now)
|
||||
launches = []
|
||||
|
||||
async def launch_run(**kwargs):
|
||||
launches.append(kwargs)
|
||||
return {"run_id": f"run-{action}", "thread_id": kwargs["thread_id"]}
|
||||
|
||||
service = _make_service(task_repo, run_repo, launch_run)
|
||||
dispatch = asyncio.create_task(service.dispatch_task(task, now=now, trigger="manual"))
|
||||
await run_repo.before_insert.wait()
|
||||
|
||||
if action == "pause":
|
||||
mutation = await task_repo.pause_with_queue_cancellation(
|
||||
task["id"],
|
||||
user_id="user-1",
|
||||
error="paused before admission",
|
||||
now=now + timedelta(seconds=1),
|
||||
)
|
||||
assert mutation == "paused"
|
||||
else:
|
||||
mutation = await task_repo.delete_with_queue_cancellation(
|
||||
task["id"],
|
||||
user_id="user-1",
|
||||
error="deleted before admission",
|
||||
now=now + timedelta(seconds=1),
|
||||
)
|
||||
assert mutation == "deleted"
|
||||
|
||||
run_repo.resume_insert.set()
|
||||
result = await dispatch
|
||||
|
||||
assert launches == []
|
||||
assert await run_repo.list_by_task(task["id"]) == []
|
||||
if action == "pause":
|
||||
assert result["outcome"] == "conflict"
|
||||
assert (await task_repo.get(task["id"], user_id="user-1"))["status"] == "paused"
|
||||
else:
|
||||
assert result["outcome"] == "not_found"
|
||||
assert await task_repo.get(task["id"], user_id="user-1") is None
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_scheduled_queue_timeout_advances_cron_without_immediate_requeue(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
task_repo = ScheduledTaskRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
admitted_at = datetime.now(UTC)
|
||||
await task_repo.create(
|
||||
task_id="task-timeout",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="Timeout",
|
||||
prompt="Prompt",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "* * * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=admitted_at - timedelta(minutes=1),
|
||||
)
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-timeout",
|
||||
task_id="task-timeout",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=admitted_at,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
|
||||
async def launch_run(**_kwargs):
|
||||
raise AssertionError("an expired queue row must not launch")
|
||||
|
||||
service = _make_service(
|
||||
task_repo,
|
||||
run_repo,
|
||||
launch_run,
|
||||
queue_timeout_seconds=60,
|
||||
)
|
||||
poll_at = admitted_at + timedelta(seconds=61)
|
||||
|
||||
await service.run_once(now=poll_at)
|
||||
|
||||
rows = await run_repo.list_by_task("task-timeout")
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["status"] == "failed"
|
||||
task = await task_repo.get("task-timeout", user_id="user-1")
|
||||
assert task is not None
|
||||
assert task["status"] == "enabled"
|
||||
assert datetime.fromisoformat(task["next_run_at"]) > poll_at
|
||||
assert task["last_error"] == "scheduled task queue wait timeout exceeded"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_manual_queue_timeout_preserves_serialized_next_run_at(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
task_repo = ScheduledTaskRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
admitted_at = datetime.now(UTC)
|
||||
next_run_at = admitted_at + timedelta(days=1)
|
||||
await task_repo.create(
|
||||
task_id="task-manual-timeout",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="Manual timeout",
|
||||
prompt="Prompt",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "0 9 * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=next_run_at,
|
||||
)
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-manual-timeout",
|
||||
task_id="task-manual-timeout",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=admitted_at,
|
||||
trigger="manual",
|
||||
status="queued",
|
||||
)
|
||||
|
||||
async def launch_run(**_kwargs):
|
||||
raise AssertionError("an expired queue row must not launch")
|
||||
|
||||
service = _make_service(
|
||||
task_repo,
|
||||
run_repo,
|
||||
launch_run,
|
||||
queue_timeout_seconds=60,
|
||||
)
|
||||
|
||||
await service.run_once(now=admitted_at + timedelta(seconds=61))
|
||||
|
||||
row = (await run_repo.list_by_task("task-manual-timeout"))[0]
|
||||
assert row["status"] == "failed"
|
||||
task = await task_repo.get("task-manual-timeout", user_id="user-1")
|
||||
assert task is not None
|
||||
assert task["status"] == "enabled"
|
||||
assert task["next_run_at"] == next_run_at.isoformat()
|
||||
assert task["last_error"] == "scheduled task queue wait timeout exceeded"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_expired_launch_claim_attaches_existing_run_instead_of_relaunching(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
task_repo = ScheduledTaskRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
durable_runs = RunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
launched_at = now + timedelta(seconds=1)
|
||||
await _seed_reuse_task(task_repo, task_id="task-attached", now=now)
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-attached",
|
||||
task_id="task-attached",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=now,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
assert (
|
||||
await run_repo.claim_queued_run(
|
||||
"task-run-attached",
|
||||
lease_owner="worker-a",
|
||||
now=now,
|
||||
lease_seconds=5,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
is not None
|
||||
)
|
||||
assert await run_repo.requeue_claimed_run(
|
||||
"task-run-attached",
|
||||
lease_owner="worker-a",
|
||||
error="earlier overlap",
|
||||
)
|
||||
assert (
|
||||
await run_repo.claim_queued_run(
|
||||
"task-run-attached",
|
||||
lease_owner="worker-a",
|
||||
now=now,
|
||||
lease_seconds=5,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
is not None
|
||||
)
|
||||
await durable_runs.put(
|
||||
"run-attached",
|
||||
thread_id="thread-1",
|
||||
user_id="user-1",
|
||||
status="running",
|
||||
metadata={
|
||||
"scheduled_task_id": "task-attached",
|
||||
"scheduled_task_run_id": "task-run-attached",
|
||||
},
|
||||
created_at=launched_at.isoformat(),
|
||||
)
|
||||
|
||||
recovered = await run_repo.recover_expired_launch_claims(
|
||||
error="launch lease expired",
|
||||
now=now + timedelta(seconds=6),
|
||||
)
|
||||
|
||||
assert recovered == 1
|
||||
row = (await run_repo.list_by_task("task-attached"))[0]
|
||||
assert row["status"] == "running"
|
||||
assert row["run_id"] == "run-attached"
|
||||
assert row["lease_owner"] is None
|
||||
assert row["started_at"] == launched_at.isoformat()
|
||||
assert row["error"] is None
|
||||
task = await task_repo.get("task-attached", user_id="user-1")
|
||||
assert task is not None
|
||||
assert task["last_run_id"] == "run-attached"
|
||||
assert task["run_count"] == 1
|
||||
|
||||
# If the original launch coroutine resumes after recovery, its normal
|
||||
# parent update must be idempotent for the same durable run id.
|
||||
await task_repo.update_after_launch(
|
||||
"task-attached",
|
||||
status="enabled",
|
||||
next_run_at=datetime.fromisoformat(task["next_run_at"]),
|
||||
last_run_at=now,
|
||||
last_run_id="run-attached",
|
||||
last_thread_id="thread-1",
|
||||
last_error=None,
|
||||
increment_run_count=True,
|
||||
)
|
||||
assert (await task_repo.get("task-attached", user_id="user-1"))["run_count"] == 1
|
||||
finally:
|
||||
await close_engine()
|
||||
@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import pytest
|
||||
@ -5,8 +6,14 @@ 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.run import RunRepository
|
||||
from deerflow.persistence.scheduled_task_runs import ActiveScheduledRunConflict, ScheduledTaskRunRepository
|
||||
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
|
||||
from deerflow.persistence.scheduled_task_runs import (
|
||||
ActiveScheduledRunConflict,
|
||||
ScheduledTaskAdmissionRejected,
|
||||
ScheduledTaskRunRepository,
|
||||
)
|
||||
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
|
||||
from deerflow.persistence.scheduled_tasks import ActiveScheduledTaskMutationConflict, ScheduledTaskRepository
|
||||
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -37,6 +44,118 @@ async def test_scheduled_task_repository_create_and_list(tmp_path):
|
||||
await close_engine()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mutable_update_rechecks_active_occurrence_at_commit_boundary(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
task_repo = ScheduledTaskRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
now = datetime.now(UTC)
|
||||
await task_repo.create(
|
||||
task_id="task-atomic-patch",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="atomic patch",
|
||||
prompt="original prompt",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "0 9 * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=None,
|
||||
)
|
||||
|
||||
# Model the router's earlier fast-path check, then admit an occurrence
|
||||
# before the actual mutation reaches the repository transaction.
|
||||
assert await task_repo.get_active_run_status("task-atomic-patch") is None
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-atomic-patch",
|
||||
task_id="task-atomic-patch",
|
||||
thread_id="thread-1",
|
||||
scheduled_for=now,
|
||||
trigger="manual",
|
||||
status="queued",
|
||||
)
|
||||
|
||||
with pytest.raises(ActiveScheduledTaskMutationConflict, match="active queued"):
|
||||
await task_repo.update(
|
||||
"task-atomic-patch",
|
||||
user_id="user-1",
|
||||
updates={"prompt": "changed after admission"},
|
||||
require_mutable=True,
|
||||
)
|
||||
|
||||
task = await task_repo.get("task-atomic-patch", user_id="user-1")
|
||||
assert task is not None
|
||||
assert task["prompt"] == "original prompt"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_atomic_update_and_admission_cannot_both_commit(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
task_repo = ScheduledTaskRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
|
||||
for index in range(5):
|
||||
task_id = f"task-update-admission-race-{index}"
|
||||
task = await task_repo.create(
|
||||
task_id=task_id,
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="atomic race",
|
||||
prompt="original prompt",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "0 9 * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=None,
|
||||
)
|
||||
|
||||
update_result, admission_result = await asyncio.gather(
|
||||
task_repo.update(
|
||||
task_id,
|
||||
user_id="user-1",
|
||||
updates={"prompt": "updated prompt"},
|
||||
require_mutable=True,
|
||||
),
|
||||
run_repo.create(
|
||||
run_record_id=f"task-run-update-admission-race-{index}",
|
||||
task_id=task_id,
|
||||
thread_id="thread-1",
|
||||
scheduled_for=datetime.now(UTC),
|
||||
trigger="manual",
|
||||
status="queued",
|
||||
coordinate_with_task=True,
|
||||
expected_task_user_id="user-1",
|
||||
expected_task_status=task["status"],
|
||||
expected_task_updated_at=task["updated_at"],
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
successes = sum(not isinstance(result, Exception) for result in (update_result, admission_result))
|
||||
assert successes == 1
|
||||
assert isinstance(update_result, ActiveScheduledTaskMutationConflict) or isinstance(
|
||||
admission_result,
|
||||
ScheduledTaskAdmissionRejected,
|
||||
)
|
||||
|
||||
current = await task_repo.get(task_id, user_id="user-1")
|
||||
assert current is not None
|
||||
active = await run_repo.get_active_run(task_id)
|
||||
assert (current["prompt"] == "updated prompt") is (active is None)
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduled_task_run_repository_records_history(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
@ -61,8 +180,7 @@ async def test_scheduled_task_run_repository_records_history(tmp_path):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_stale_active_runs_fails_orphaned_runs(tmp_path):
|
||||
"""Runs stuck in queued/running after a process crash are swept to interrupted."""
|
||||
async def test_mark_stale_active_runs_preserves_queue_and_interrupts_live_run(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
@ -100,11 +218,11 @@ async def test_mark_stale_active_runs_fails_orphaned_runs(tmp_path):
|
||||
)
|
||||
|
||||
swept = await repo.mark_stale_active_runs(error="interrupted: gateway restarted")
|
||||
assert swept == 2
|
||||
assert swept == 1
|
||||
|
||||
by_id = {entry["id"]: entry for entry in await repo.list_by_task("task-1")}
|
||||
by_id.update({entry["id"]: entry for entry in await repo.list_by_task("task-2")})
|
||||
assert by_id["task-run-queued"]["status"] == "interrupted"
|
||||
assert by_id["task-run-queued"]["status"] == "queued"
|
||||
assert by_id["task-run-running"]["status"] == "interrupted"
|
||||
assert by_id["task-run-success"]["status"] == "success"
|
||||
|
||||
@ -184,6 +302,189 @@ async def test_lease_aware_recovery_preserves_live_peer_and_reclaims_expired_pee
|
||||
await close_engine()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_live_launch_repairs_bookkeeping_before_releasing_claim(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)
|
||||
task_run_repo = ScheduledTaskRunRepository(sf, run_repository=durable_run_repo)
|
||||
now = datetime(2026, 8, 21, 1, 30, tzinfo=UTC)
|
||||
launched_at = now + timedelta(seconds=1)
|
||||
|
||||
await task_repo.create(
|
||||
task_id="task-live-launch",
|
||||
user_id="user-1",
|
||||
thread_id="thread-live-launch",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="live launch",
|
||||
prompt="p",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "* * * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=None,
|
||||
)
|
||||
await task_run_repo.create(
|
||||
run_record_id="task-run-live-launch",
|
||||
task_id="task-live-launch",
|
||||
thread_id="thread-live-launch",
|
||||
scheduled_for=now,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
assert (
|
||||
await task_run_repo.claim_queued_run(
|
||||
"task-run-live-launch",
|
||||
lease_owner="pod-launcher",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
is not None
|
||||
)
|
||||
assert await task_run_repo.requeue_claimed_run(
|
||||
"task-run-live-launch",
|
||||
lease_owner="pod-launcher",
|
||||
error="earlier overlap",
|
||||
)
|
||||
assert (
|
||||
await task_run_repo.claim_queued_run(
|
||||
"task-run-live-launch",
|
||||
lease_owner="pod-launcher",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
is not None
|
||||
)
|
||||
await durable_run_repo.put(
|
||||
"run-live-launch",
|
||||
thread_id="thread-live-launch",
|
||||
user_id="user-1",
|
||||
status="running",
|
||||
metadata={
|
||||
"scheduled_task_id": "task-live-launch",
|
||||
"scheduled_task_run_id": "task-run-live-launch",
|
||||
},
|
||||
created_at=launched_at.isoformat(),
|
||||
owner_worker_id="pod-launcher",
|
||||
lease_expires_at=(now + timedelta(seconds=120)).isoformat(),
|
||||
)
|
||||
|
||||
assert await task_run_repo.reconcile_active_runs(error="lease expired", now=now) == 0
|
||||
row = (await task_run_repo.list_by_task("task-live-launch"))[0]
|
||||
assert row["status"] == "running"
|
||||
assert row["run_id"] == "run-live-launch"
|
||||
assert row["lease_owner"] is None
|
||||
assert row["started_at"] == launched_at.isoformat()
|
||||
assert row["error"] is None
|
||||
|
||||
# The original launcher is now fenced because reconciliation released
|
||||
# its short claim, but the stable row already contains all bookkeeping.
|
||||
assert not await task_run_repo.update_status(
|
||||
"task-run-live-launch",
|
||||
status="running",
|
||||
run_id="run-live-launch",
|
||||
started_at=launched_at,
|
||||
protect_terminal=True,
|
||||
expected_lease_owner="pod-launcher",
|
||||
)
|
||||
row = (await task_run_repo.list_by_task("task-live-launch"))[0]
|
||||
assert row["started_at"] == launched_at.isoformat()
|
||||
assert row["error"] is None
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reconcile_locks_task_before_its_active_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)
|
||||
|
||||
seed_repo = ScheduledTaskRunRepository(sf)
|
||||
# Insert in reverse lexical order: every reconciler must still acquire
|
||||
# task/run pairs in one deterministic global order.
|
||||
for suffix in ("z", "a"):
|
||||
await task_repo.create(
|
||||
task_id=f"task-lock-order-{suffix}",
|
||||
user_id="user-1",
|
||||
thread_id=None,
|
||||
context_mode="fresh_thread_per_run",
|
||||
assistant_id="lead_agent",
|
||||
title="lock order",
|
||||
prompt="p",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "* * * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=None,
|
||||
)
|
||||
await seed_repo.create(
|
||||
run_record_id=f"task-run-lock-order-{suffix}",
|
||||
task_id=f"task-lock-order-{suffix}",
|
||||
thread_id=f"thread-lock-order-{suffix}",
|
||||
scheduled_for=now,
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
assert (
|
||||
await seed_repo.claim_queued_run(
|
||||
f"task-run-lock-order-{suffix}",
|
||||
lease_owner="pod-a",
|
||||
now=now,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
is not None
|
||||
)
|
||||
|
||||
lock_order = []
|
||||
|
||||
class RecordingSession:
|
||||
def __init__(self, session):
|
||||
self._session = session
|
||||
|
||||
async def __aenter__(self):
|
||||
await self._session.__aenter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return await self._session.__aexit__(*args)
|
||||
|
||||
async def get(self, entity, ident, **kwargs):
|
||||
if kwargs.get("with_for_update"):
|
||||
lock_order.append((entity, ident))
|
||||
return await self._session.get(entity, ident, **kwargs)
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._session, name)
|
||||
|
||||
def recording_session_factory():
|
||||
return RecordingSession(sf())
|
||||
|
||||
reconcile_repo = ScheduledTaskRunRepository(
|
||||
recording_session_factory,
|
||||
run_repository=durable_run_repo,
|
||||
)
|
||||
|
||||
assert await reconcile_repo.reconcile_active_runs(error="lease expired", now=now) == 0
|
||||
assert lock_order == [
|
||||
(ScheduledTaskRow, "task-lock-order-a"),
|
||||
(ScheduledTaskRunRow, "task-run-lock-order-a"),
|
||||
(ScheduledTaskRow, "task-lock-order-z"),
|
||||
(ScheduledTaskRunRow, "task-run-lock-order-z"),
|
||||
]
|
||||
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)))
|
||||
@ -219,8 +520,8 @@ async def test_lease_aware_recovery_preserves_queued_dispatch_until_lease_expire
|
||||
|
||||
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"
|
||||
assert await task_run_repo.reconcile_active_runs(error="restart", now=now + timedelta(seconds=121)) == 0
|
||||
assert (await task_run_repo.list_by_task("task-queued"))[0]["status"] == "queued"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
@ -289,21 +590,53 @@ async def test_update_status_protect_terminal_keeps_completion_result(tmp_path):
|
||||
trigger="scheduled",
|
||||
status="queued",
|
||||
)
|
||||
claimed_at = datetime(2026, 7, 2, 1, 0, tzinfo=UTC)
|
||||
claimed = await repo.claim_queued_run(
|
||||
"task-run-race",
|
||||
lease_owner="worker-a",
|
||||
now=claimed_at,
|
||||
lease_seconds=120,
|
||||
global_max_concurrent_runs=3,
|
||||
)
|
||||
assert claimed is not None
|
||||
# Completion hook wins the race and commits the terminal state first.
|
||||
await repo.update_status("task-run-race", status="failed", run_id="run-1", error="boom", finished_at=datetime(2026, 7, 2, 1, 1, tzinfo=UTC))
|
||||
# Late launch-path write: keeps terminal status/error, backfills started_at.
|
||||
await repo.update_status("task-run-race", status="running", run_id="run-1", started_at=datetime(2026, 7, 2, 1, 0, tzinfo=UTC), protect_terminal=True)
|
||||
# Late launch-path write: completion cleared the launch lease, but the
|
||||
# matching run id proves this is the same launch and permits the missing
|
||||
# started_at backfill without weakening fencing for another run.
|
||||
updated = await repo.update_status(
|
||||
"task-run-race",
|
||||
status="running",
|
||||
run_id="run-1",
|
||||
started_at=claimed_at,
|
||||
protect_terminal=True,
|
||||
expected_lease_owner="worker-a",
|
||||
)
|
||||
|
||||
entry = (await repo.list_by_task("task-1"))[0]
|
||||
assert updated is True
|
||||
assert entry["status"] == "failed"
|
||||
assert entry["error"] == "boom"
|
||||
assert entry["started_at"] is not None
|
||||
|
||||
stale = await repo.update_status(
|
||||
"task-run-race",
|
||||
status="running",
|
||||
run_id="run-stale",
|
||||
started_at=claimed_at - timedelta(minutes=1),
|
||||
protect_terminal=True,
|
||||
expected_lease_owner="worker-stale",
|
||||
)
|
||||
assert stale is False
|
||||
entry = (await repo.list_by_task("task-1"))[0]
|
||||
assert entry["run_id"] == "run-1"
|
||||
assert entry["started_at"] == claimed_at.isoformat()
|
||||
|
||||
await close_engine()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_has_active_runs_sees_only_queued_and_running(tmp_path):
|
||||
async def test_has_active_runs_sees_all_nonterminal_queue_states(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
@ -319,6 +652,8 @@ async def test_has_active_runs_sees_only_queued_and_running(tmp_path):
|
||||
status="running",
|
||||
)
|
||||
assert await repo.has_active_runs("task-1") is True
|
||||
await repo.update_status("task-run-active", status="launching")
|
||||
assert await repo.has_active_runs("task-1") is True
|
||||
await repo.update_status("task-run-active", status="success", run_id="run-1")
|
||||
assert await repo.has_active_runs("task-1") is False
|
||||
|
||||
@ -839,40 +1174,3 @@ async def test_update_after_launch_rejects_stale_lease_owner(tmp_path, caplog):
|
||||
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()
|
||||
|
||||
@ -1,10 +1,16 @@
|
||||
import asyncio
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from app.gateway.routers import scheduled_tasks
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
|
||||
from deerflow.persistence.scheduled_task_runs import ScheduledTaskRunRepository
|
||||
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@ -32,10 +38,27 @@ def test_scheduled_task_models_reject_invalid_thread_ids(model, thread_id):
|
||||
model(**kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status", "offers_pause_cancellation"),
|
||||
[
|
||||
("queued", True),
|
||||
("launching", False),
|
||||
("running", False),
|
||||
],
|
||||
)
|
||||
def test_active_occurrence_conflict_detail_only_offers_pause_for_queued(status, offers_pause_cancellation):
|
||||
detail = scheduled_tasks._active_occurrence_conflict_detail(status)
|
||||
|
||||
assert f"active {status} occurrence" in detail
|
||||
assert ("cancel the queued occurrence by pausing the task" in detail) is offers_pause_cancellation
|
||||
|
||||
|
||||
class _Repo:
|
||||
def __init__(self) -> None:
|
||||
self.created = []
|
||||
self.items = {}
|
||||
self.active_status = None
|
||||
self.cancelled_queue = False
|
||||
|
||||
async def list_by_user(self, user_id: str):
|
||||
return [item for item in self.items.values() if item["user_id"] == user_id]
|
||||
@ -67,13 +90,40 @@ class _Repo:
|
||||
return None
|
||||
return item
|
||||
|
||||
async def update(self, task_id: str, *, user_id: str, updates):
|
||||
async def update(self, task_id: str, *, user_id: str, updates, require_mutable: bool = False):
|
||||
item = await self.get(task_id, user_id=user_id)
|
||||
if item is None:
|
||||
return None
|
||||
item.update(updates)
|
||||
return item
|
||||
|
||||
async def get_active_run_status(self, task_id: str):
|
||||
return self.active_status
|
||||
|
||||
async def pause_with_queue_cancellation(self, task_id: str, *, user_id: str, **_kwargs):
|
||||
item = await self.get(task_id, user_id=user_id)
|
||||
if item is None:
|
||||
return "not_found"
|
||||
if self.active_status in {"launching", "running"}:
|
||||
return "executing"
|
||||
if self.active_status == "queued":
|
||||
self.cancelled_queue = True
|
||||
self.active_status = None
|
||||
item["status"] = "paused"
|
||||
return "paused"
|
||||
|
||||
async def delete_with_queue_cancellation(self, task_id: str, *, user_id: str, **_kwargs):
|
||||
item = await self.get(task_id, user_id=user_id)
|
||||
if item is None:
|
||||
return "not_found"
|
||||
if self.active_status in {"launching", "running"}:
|
||||
return "executing"
|
||||
if self.active_status == "queued":
|
||||
self.cancelled_queue = True
|
||||
self.active_status = None
|
||||
self.items.pop(task_id, None)
|
||||
return "deleted"
|
||||
|
||||
async def delete(self, task_id: str, *, user_id: str):
|
||||
item = await self.get(task_id, user_id=user_id)
|
||||
if item is None:
|
||||
@ -324,6 +374,83 @@ async def test_update_scheduled_task_writes_repo():
|
||||
assert result["title"] == "Updated title"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_rechecks_atomic_mutability_after_router_precheck(tmp_path):
|
||||
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
||||
old_repo = scheduled_tasks.get_scheduled_task_repo
|
||||
old_thread_store = scheduled_tasks.get_thread_store
|
||||
old_config = scheduled_tasks.get_config
|
||||
old_user = scheduled_tasks.get_optional_user_from_request
|
||||
try:
|
||||
sf = get_session_factory()
|
||||
assert sf is not None
|
||||
|
||||
class PrecheckBarrierRepository(ScheduledTaskRepository):
|
||||
def __init__(self, session_factory):
|
||||
super().__init__(session_factory)
|
||||
self.prechecked = asyncio.Event()
|
||||
self.resume = asyncio.Event()
|
||||
|
||||
async def get_active_run_status(self, task_id: str):
|
||||
status = await super().get_active_run_status(task_id)
|
||||
if status is None:
|
||||
self.prechecked.set()
|
||||
await self.resume.wait()
|
||||
return status
|
||||
|
||||
repo = PrecheckBarrierRepository(sf)
|
||||
run_repo = ScheduledTaskRunRepository(sf)
|
||||
task = await repo.create(
|
||||
task_id="task-router-atomic-patch",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="Atomic patch",
|
||||
prompt="original prompt",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "0 9 * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=None,
|
||||
)
|
||||
scheduled_tasks.get_scheduled_task_repo = lambda _request: repo
|
||||
scheduled_tasks.get_thread_store = lambda _request: SimpleNamespace(check_access=AsyncMock(return_value=True))
|
||||
scheduled_tasks.get_config = lambda: _Config()
|
||||
scheduled_tasks.get_optional_user_from_request = AsyncMock(return_value=SimpleNamespace(id="user-1"))
|
||||
|
||||
patch_call = asyncio.create_task(
|
||||
scheduled_tasks.update_scheduled_task.__wrapped__(
|
||||
task_id=task["id"],
|
||||
request=SimpleNamespace(),
|
||||
body=scheduled_tasks.ScheduledTaskUpdateRequest(prompt="changed after admission"),
|
||||
)
|
||||
)
|
||||
await repo.prechecked.wait()
|
||||
await run_repo.create(
|
||||
run_record_id="task-run-router-atomic-patch",
|
||||
task_id=task["id"],
|
||||
thread_id="thread-1",
|
||||
scheduled_for=datetime.now(UTC),
|
||||
trigger="manual",
|
||||
status="queued",
|
||||
)
|
||||
repo.resume.set()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await patch_call
|
||||
assert exc_info.value.status_code == 409
|
||||
assert "active queued occurrence" in exc_info.value.detail
|
||||
current = await repo.get(task["id"], user_id="user-1")
|
||||
assert current is not None
|
||||
assert current["prompt"] == "original prompt"
|
||||
finally:
|
||||
scheduled_tasks.get_scheduled_task_repo = old_repo
|
||||
scheduled_tasks.get_thread_store = old_thread_store
|
||||
scheduled_tasks.get_config = old_config
|
||||
scheduled_tasks.get_optional_user_from_request = old_user
|
||||
await close_engine()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_scheduled_task_deletes_repo_row():
|
||||
repo = _Repo()
|
||||
@ -403,6 +530,81 @@ async def test_pause_and_resume_scheduled_task_update_status():
|
||||
assert resumed["status"] == "enabled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pause_cancels_waiting_occurrence_before_pausing_task():
|
||||
repo = _Repo()
|
||||
task = await repo.create(
|
||||
task_id="task-queued",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="Queued task",
|
||||
prompt="Prompt",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "0 9 * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=None,
|
||||
)
|
||||
repo.active_status = "queued"
|
||||
request = SimpleNamespace()
|
||||
user = SimpleNamespace(id="user-1")
|
||||
|
||||
old_repo = scheduled_tasks.get_scheduled_task_repo
|
||||
old_user = scheduled_tasks.get_optional_user_from_request
|
||||
try:
|
||||
scheduled_tasks.get_scheduled_task_repo = lambda _request: repo
|
||||
scheduled_tasks.get_optional_user_from_request = AsyncMock(return_value=user)
|
||||
result = await scheduled_tasks.pause_scheduled_task.__wrapped__(
|
||||
task_id=task["id"],
|
||||
request=request,
|
||||
)
|
||||
finally:
|
||||
scheduled_tasks.get_scheduled_task_repo = old_repo
|
||||
scheduled_tasks.get_optional_user_from_request = old_user
|
||||
|
||||
assert result["status"] == "paused"
|
||||
assert repo.cancelled_queue is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_rejects_occurrence_that_has_started_launching():
|
||||
repo = _Repo()
|
||||
task = await repo.create(
|
||||
task_id="task-launching",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="Launching task",
|
||||
prompt="Prompt",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "0 9 * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=None,
|
||||
)
|
||||
repo.active_status = "launching"
|
||||
request = SimpleNamespace()
|
||||
user = SimpleNamespace(id="user-1")
|
||||
|
||||
old_repo = scheduled_tasks.get_scheduled_task_repo
|
||||
old_user = scheduled_tasks.get_optional_user_from_request
|
||||
try:
|
||||
scheduled_tasks.get_scheduled_task_repo = lambda _request: repo
|
||||
scheduled_tasks.get_optional_user_from_request = AsyncMock(return_value=user)
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await scheduled_tasks.delete_scheduled_task.__wrapped__(
|
||||
task_id=task["id"],
|
||||
request=request,
|
||||
)
|
||||
finally:
|
||||
scheduled_tasks.get_scheduled_task_repo = old_repo
|
||||
scheduled_tasks.get_optional_user_from_request = old_user
|
||||
|
||||
assert "launching or running" in str(exc_info.value)
|
||||
assert task["id"] in repo.items
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pause_rejects_running_task():
|
||||
repo = _Repo()
|
||||
@ -488,6 +690,50 @@ async def test_update_rejects_running_task():
|
||||
assert "currently running" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_rejects_queued_task_definition_until_occurrence_finishes():
|
||||
repo = _Repo()
|
||||
task = await repo.create(
|
||||
task_id="task-queued",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
context_mode="reuse_thread",
|
||||
assistant_id="lead_agent",
|
||||
title="Queued task",
|
||||
prompt="Original prompt",
|
||||
schedule_type="cron",
|
||||
schedule_spec={"cron": "0 9 * * *"},
|
||||
timezone="UTC",
|
||||
next_run_at=None,
|
||||
)
|
||||
repo.active_status = "queued"
|
||||
request = SimpleNamespace()
|
||||
user = SimpleNamespace(id="user-1")
|
||||
config = _Config()
|
||||
|
||||
old_repo = scheduled_tasks.get_scheduled_task_repo
|
||||
old_config = scheduled_tasks.get_config
|
||||
old_user = scheduled_tasks.get_optional_user_from_request
|
||||
try:
|
||||
scheduled_tasks.get_scheduled_task_repo = lambda _request: repo
|
||||
scheduled_tasks.get_config = lambda: config
|
||||
scheduled_tasks.get_optional_user_from_request = AsyncMock(return_value=user)
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await scheduled_tasks.update_scheduled_task.__wrapped__(
|
||||
task_id=task["id"],
|
||||
request=request,
|
||||
body=scheduled_tasks.ScheduledTaskUpdateRequest(prompt="Changed while queued"),
|
||||
)
|
||||
finally:
|
||||
scheduled_tasks.get_scheduled_task_repo = old_repo
|
||||
scheduled_tasks.get_config = old_config
|
||||
scheduled_tasks.get_optional_user_from_request = old_user
|
||||
|
||||
assert "active queued occurrence" in str(exc_info.value)
|
||||
assert repo.items[task["id"]]["prompt"] == "Original prompt"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_thread_scheduled_tasks_filters_by_thread_id():
|
||||
repo = _Repo()
|
||||
|
||||
@ -14,6 +14,7 @@ class DummyTaskRepo:
|
||||
self.rows = rows
|
||||
self.claimed = False
|
||||
self.updated = None
|
||||
self.release_calls = []
|
||||
self.cancelled_stuck_once = None
|
||||
self.reconciled_stuck_once = None
|
||||
|
||||
@ -28,6 +29,13 @@ class DummyTaskRepo:
|
||||
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 release_queued_admission_lease(self, task_id):
|
||||
return False
|
||||
|
||||
async def release_dispatch_lease(self, task_id, **kwargs):
|
||||
self.release_calls.append((task_id, kwargs))
|
||||
return True
|
||||
|
||||
async def claim_due_tasks(self, **_kwargs):
|
||||
if self.claimed:
|
||||
return []
|
||||
@ -41,6 +49,10 @@ class DummyTaskRepo:
|
||||
row = next((item for item in self.rows if item["id"] == task_id and item["user_id"] == user_id), None)
|
||||
return dict(row) if row is not None else None
|
||||
|
||||
async def get_internal(self, task_id: str):
|
||||
row = next((item for item in self.rows if item["id"] == task_id), None)
|
||||
return dict(row) if row is not None else None
|
||||
|
||||
async def update(self, task_id: str, *, user_id: str, updates):
|
||||
row = next((item for item in self.rows if item["id"] == task_id and item["user_id"] == user_id), None)
|
||||
if row is None:
|
||||
@ -62,12 +74,49 @@ class DummyRunRepo:
|
||||
async def count_active_runs(self):
|
||||
return self.active_count
|
||||
|
||||
async def list_queued_runs(self, *, limit):
|
||||
return []
|
||||
|
||||
async def expire_queued_runs(self, **_kwargs):
|
||||
return []
|
||||
|
||||
async def recover_expired_launch_claims(self, **_kwargs):
|
||||
return 0
|
||||
|
||||
async def get_active_run(self, task_id):
|
||||
if not self.active:
|
||||
return None
|
||||
return {
|
||||
"id": "task-run-active",
|
||||
"task_id": task_id,
|
||||
"thread_id": "thread-active",
|
||||
"status": "running",
|
||||
}
|
||||
|
||||
async def claim_queued_run(self, run_record_id, *, global_max_concurrent_runs, **_kwargs):
|
||||
if self.active_count >= global_max_concurrent_runs:
|
||||
return None
|
||||
return {"id": run_record_id, "status": "launching"}
|
||||
|
||||
async def requeue_claimed_run(self, run_record_id, **kwargs):
|
||||
self.updated.append((run_record_id, {"status": "queued", **kwargs}))
|
||||
return True
|
||||
|
||||
async def create(self, **kwargs):
|
||||
self.created = kwargs
|
||||
return {"id": kwargs["run_record_id"]}
|
||||
|
||||
async def update_status(self, run_record_id, **kwargs):
|
||||
self.updated.append((run_record_id, kwargs))
|
||||
return True
|
||||
|
||||
async def reconcile_launched_run(self, run_record_id, **kwargs):
|
||||
self.updated.append((run_record_id, {"reconciled": True, **kwargs}))
|
||||
return True
|
||||
|
||||
async def fail_launching_run(self, run_record_id, **kwargs):
|
||||
self.updated.append((run_record_id, {"status": "failed", **kwargs}))
|
||||
return True
|
||||
|
||||
async def has_active_runs(self, task_id):
|
||||
return self.active
|
||||
@ -207,7 +256,7 @@ async def test_fresh_thread_per_run_creates_new_execution_thread():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduled_overlap_conflict_is_recorded_as_skip():
|
||||
async def test_scheduled_overlap_conflict_is_kept_in_queue():
|
||||
async def fake_launch(**_kwargs):
|
||||
raise ConflictError("Thread thread-1 already has an active run")
|
||||
|
||||
@ -224,7 +273,7 @@ async def test_scheduled_overlap_conflict_is_recorded_as_skip():
|
||||
"schedule_spec": {"cron": "0 9 * * *"},
|
||||
"timezone": "UTC",
|
||||
"status": "running",
|
||||
"overlap_policy": "skip",
|
||||
"overlap_policy": "enqueue",
|
||||
"last_run_id": "run-old",
|
||||
"last_thread_id": "thread-1",
|
||||
"last_run_at": "2026-07-01T00:00:00+00:00",
|
||||
@ -247,13 +296,14 @@ async def test_scheduled_overlap_conflict_is_recorded_as_skip():
|
||||
trigger="scheduled",
|
||||
)
|
||||
|
||||
assert result["outcome"] == "skipped"
|
||||
assert run_repo.updated[-1][1]["status"] == "skipped"
|
||||
assert task_repo.updated[1]["status"] == "enabled"
|
||||
assert result["outcome"] == "queued"
|
||||
assert run_repo.created["status"] == "queued"
|
||||
assert run_repo.updated[-1][1]["status"] == "queued"
|
||||
assert task_repo.updated is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manual_overlap_conflict_returns_conflict():
|
||||
async def test_manual_overlap_conflict_is_kept_in_queue():
|
||||
async def fake_launch(**_kwargs):
|
||||
raise ConflictError("Thread thread-1 already has an active run")
|
||||
|
||||
@ -270,7 +320,7 @@ async def test_manual_overlap_conflict_returns_conflict():
|
||||
"schedule_spec": {"cron": "0 9 * * *"},
|
||||
"timezone": "UTC",
|
||||
"status": "enabled",
|
||||
"overlap_policy": "skip",
|
||||
"overlap_policy": "enqueue",
|
||||
}
|
||||
]
|
||||
)
|
||||
@ -290,8 +340,9 @@ async def test_manual_overlap_conflict_returns_conflict():
|
||||
trigger="manual",
|
||||
)
|
||||
|
||||
assert result["outcome"] == "conflict"
|
||||
assert run_repo.updated[-1][1]["status"] == "failed"
|
||||
assert result["outcome"] == "queued"
|
||||
assert run_repo.updated[-1][1]["status"] == "queued"
|
||||
assert task_repo.release_calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -545,7 +596,7 @@ async def test_interrupted_cron_run_keeps_task_enabled():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skip_policy_applies_to_fresh_thread_runs():
|
||||
async def test_existing_running_occurrence_blocks_duplicate_fresh_thread_run():
|
||||
launched = []
|
||||
|
||||
async def fake_launch(**kwargs):
|
||||
@ -553,7 +604,7 @@ async def test_skip_policy_applies_to_fresh_thread_runs():
|
||||
return {"run_id": "run-9", "thread_id": kwargs["thread_id"]}
|
||||
|
||||
row = _once_task_row(task_id="task-9")
|
||||
row.update({"schedule_type": "cron", "schedule_spec": {"cron": "* * * * *"}, "status": "running", "overlap_policy": "skip"})
|
||||
row.update({"schedule_type": "cron", "schedule_spec": {"cron": "* * * * *"}, "status": "running", "overlap_policy": "enqueue"})
|
||||
task_repo = DummyTaskRepo([row])
|
||||
run_repo = DummyRunRepo(active=True)
|
||||
service = ScheduledTaskService(
|
||||
@ -567,16 +618,9 @@ async def test_skip_policy_applies_to_fresh_thread_runs():
|
||||
|
||||
result = await service.dispatch_task(row, now=datetime.now(UTC), trigger="scheduled")
|
||||
|
||||
assert result["outcome"] == "skipped"
|
||||
assert result["outcome"] == "conflict"
|
||||
assert launched == []
|
||||
# The skip tombstone is created directly as terminal "skipped" (not the
|
||||
# transient "queued" the launch path uses): a queued row is active and would
|
||||
# itself trip the uq_scheduled_task_run_active partial unique index against
|
||||
# the pre-existing run still holding the task's single active slot.
|
||||
assert run_repo.created["status"] == "skipped"
|
||||
assert run_repo.updated[-1][1]["status"] == "skipped"
|
||||
assert task_repo.updated[1]["status"] == "enabled"
|
||||
assert task_repo.updated[1]["increment_run_count"] is False
|
||||
assert run_repo.created is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -628,7 +672,7 @@ async def test_manual_trigger_with_active_run_returns_conflict_without_launching
|
||||
return {"run_id": "run-x", "thread_id": kwargs["thread_id"]}
|
||||
|
||||
row = _once_task_row(task_id="task-manual-busy")
|
||||
row.update({"schedule_type": "cron", "schedule_spec": {"cron": "* * * * *"}, "status": "enabled", "overlap_policy": "skip"})
|
||||
row.update({"schedule_type": "cron", "schedule_spec": {"cron": "* * * * *"}, "status": "enabled", "overlap_policy": "enqueue"})
|
||||
task_repo = DummyTaskRepo([row])
|
||||
run_repo = DummyRunRepo(active=True)
|
||||
service = ScheduledTaskService(
|
||||
@ -650,7 +694,7 @@ async def test_manual_trigger_with_active_run_returns_conflict_without_launching
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_once_claims_only_into_remaining_global_budget():
|
||||
async def test_run_once_admits_due_occurrences_independently_of_execution_budget():
|
||||
claim_limits = []
|
||||
|
||||
class BudgetTaskRepo(DummyTaskRepo):
|
||||
@ -663,12 +707,11 @@ async def test_run_once_claims_only_into_remaining_global_budget():
|
||||
service = _make_service(task_repo, run_repo)
|
||||
|
||||
await service.run_once(now=datetime.now(UTC))
|
||||
assert claim_limits == [1]
|
||||
assert claim_limits == [3]
|
||||
|
||||
run_repo.active_count = 3
|
||||
await service.run_once(now=datetime.now(UTC))
|
||||
# Budget exhausted: no claim at all this cycle.
|
||||
assert claim_limits == [1]
|
||||
assert claim_limits == [3, 3]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -709,7 +752,7 @@ class _StatefulRunRepo:
|
||||
the partial unique index ``uq_scheduled_task_run_active``.
|
||||
"""
|
||||
|
||||
_ACTIVE = {"queued", "running"}
|
||||
_ACTIVE = {"queued", "launching", "running"}
|
||||
|
||||
def __init__(self, *, fail_first_update: bool = False, fail_updates: int = 0) -> None:
|
||||
self.created: list[dict] = []
|
||||
@ -719,18 +762,47 @@ class _StatefulRunRepo:
|
||||
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)
|
||||
return sum(1 for row in self.rows.values() if row["status"] in {"launching", "running"})
|
||||
|
||||
async def list_queued_runs(self, *, limit: int) -> list[dict]:
|
||||
return []
|
||||
|
||||
async def expire_queued_runs(self, **_kwargs) -> list[dict]:
|
||||
return []
|
||||
|
||||
async def create(self, **kwargs) -> dict:
|
||||
self.created.append(kwargs)
|
||||
self.rows[kwargs["run_record_id"]] = {
|
||||
"id": kwargs["run_record_id"],
|
||||
"task_id": kwargs["task_id"],
|
||||
"thread_id": kwargs["thread_id"],
|
||||
"trigger": kwargs["trigger"],
|
||||
"status": kwargs["status"],
|
||||
"run_id": None,
|
||||
}
|
||||
return {"id": kwargs["run_record_id"]}
|
||||
|
||||
async def update_status(self, run_record_id: str, **kwargs) -> None:
|
||||
async def get_active_run(self, task_id: str) -> dict | None:
|
||||
return next(
|
||||
(dict(row) for row in self.rows.values() if row["task_id"] == task_id and row["status"] in self._ACTIVE),
|
||||
None,
|
||||
)
|
||||
|
||||
async def claim_queued_run(self, run_record_id: str, **_kwargs) -> dict | None:
|
||||
row = self.rows.get(run_record_id)
|
||||
if row is None or row["status"] != "queued":
|
||||
return None
|
||||
row["status"] = "launching"
|
||||
return dict(row)
|
||||
|
||||
async def requeue_claimed_run(self, run_record_id: str, **_kwargs) -> bool:
|
||||
row = self.rows.get(run_record_id)
|
||||
if row is None or row["status"] != "launching":
|
||||
return False
|
||||
row["status"] = "queued"
|
||||
return True
|
||||
|
||||
async def update_status(self, run_record_id: str, **kwargs) -> bool:
|
||||
self.updates.append((run_record_id, kwargs))
|
||||
if self._updates_raised < self._fail_updates:
|
||||
# The launch-path queued->running write fails AFTER _launch_run has
|
||||
@ -740,11 +812,27 @@ class _StatefulRunRepo:
|
||||
raise RuntimeError("simulated transient DB error on queued->running write")
|
||||
row = self.rows.get(run_record_id)
|
||||
if row is None:
|
||||
return
|
||||
return False
|
||||
if "status" in kwargs:
|
||||
row["status"] = kwargs["status"]
|
||||
if kwargs.get("run_id") is not None:
|
||||
row["run_id"] = kwargs["run_id"]
|
||||
return True
|
||||
|
||||
async def reconcile_launched_run(self, run_record_id: str, **kwargs) -> bool:
|
||||
row = self.rows.get(run_record_id)
|
||||
if row is None:
|
||||
return False
|
||||
row["status"] = "running"
|
||||
row["run_id"] = kwargs["run_id"]
|
||||
return True
|
||||
|
||||
async def fail_launching_run(self, run_record_id: str, **kwargs) -> bool:
|
||||
row = self.rows.get(run_record_id)
|
||||
if row is None or row["status"] != "launching":
|
||||
return False
|
||||
row["status"] = "failed"
|
||||
return True
|
||||
|
||||
async def has_active_runs(self, task_id: str) -> bool:
|
||||
return any(row["task_id"] == task_id and row["status"] in self._ACTIVE for row in self.rows.values())
|
||||
@ -784,7 +872,7 @@ async def test_post_launch_bookkeeping_failure_does_not_release_active_slot():
|
||||
"schedule_spec": {"cron": "*/5 * * * *"},
|
||||
"timezone": "UTC",
|
||||
"status": "enabled",
|
||||
"overlap_policy": "skip",
|
||||
"overlap_policy": "enqueue",
|
||||
}
|
||||
]
|
||||
)
|
||||
@ -812,7 +900,7 @@ async def test_post_launch_bookkeeping_failure_does_not_release_active_slot():
|
||||
# launch a duplicate. On main (bug) this would launch run-2 here.
|
||||
second = await service.dispatch_task(task, now=now, trigger="scheduled")
|
||||
assert len(launched) == 1, launched
|
||||
assert second["outcome"] in {"skipped", "conflict"}, second
|
||||
assert second["outcome"] == "conflict", second
|
||||
|
||||
# The launched run_id is retained on the task-run row (status "running",
|
||||
# not "failed") so reconciliation / cancellation can still reach it.
|
||||
@ -858,7 +946,7 @@ async def test_both_post_launch_association_writes_can_fail_without_releasing_sl
|
||||
"schedule_spec": {"cron": "*/5 * * * *"},
|
||||
"timezone": "UTC",
|
||||
"status": "enabled",
|
||||
"overlap_policy": "skip",
|
||||
"overlap_policy": "enqueue",
|
||||
}
|
||||
task_repo = FailingTaskRepo([task])
|
||||
run_repo = _StatefulRunRepo(fail_updates=2)
|
||||
@ -875,15 +963,13 @@ async def test_both_post_launch_association_writes_can_fail_without_releasing_sl
|
||||
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,
|
||||
}
|
||||
assert run_repo.rows[first_row_id]["task_id"] == "task-double-failure"
|
||||
assert run_repo.rows[first_row_id]["status"] == "launching"
|
||||
assert run_repo.rows[first_row_id]["run_id"] is None
|
||||
|
||||
second = await service.dispatch_task(dict(task), now=now, trigger="scheduled")
|
||||
assert len(launched) == 1
|
||||
assert second["outcome"] == "skipped"
|
||||
assert second["outcome"] == "conflict"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -912,7 +998,7 @@ async def test_pre_launch_failure_still_releases_active_slot():
|
||||
"schedule_spec": {"cron": "*/5 * * * *"},
|
||||
"timezone": "UTC",
|
||||
"status": "enabled",
|
||||
"overlap_policy": "skip",
|
||||
"overlap_policy": "enqueue",
|
||||
}
|
||||
]
|
||||
)
|
||||
@ -970,7 +1056,7 @@ async def test_malformed_launch_result_still_retains_active_slot():
|
||||
"schedule_spec": {"cron": "*/5 * * * *"},
|
||||
"timezone": "UTC",
|
||||
"status": "enabled",
|
||||
"overlap_policy": "skip",
|
||||
"overlap_policy": "enqueue",
|
||||
}
|
||||
]
|
||||
)
|
||||
@ -997,15 +1083,14 @@ async def test_malformed_launch_result_still_retains_active_slot():
|
||||
# an active status, NOT "failed") and NOT launch a duplicate.
|
||||
second = await service.dispatch_task(task, now=now, trigger="scheduled")
|
||||
assert len(launched) == 1, launched
|
||||
assert second["outcome"] in {"skipped", "conflict"}, second
|
||||
assert second["outcome"] == "conflict", second
|
||||
|
||||
first_row_id = run_repo.created[0]["run_record_id"]
|
||||
assert run_repo.rows[first_row_id]["status"] == "running"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manual_trigger_rejected_when_global_budget_exhausted():
|
||||
"""Manual trigger must return a conflict when max_concurrent_runs is reached."""
|
||||
async def test_manual_trigger_is_queued_when_global_budget_exhausted():
|
||||
launched = []
|
||||
|
||||
async def fake_launch(**kwargs):
|
||||
@ -1013,7 +1098,7 @@ async def test_manual_trigger_rejected_when_global_budget_exhausted():
|
||||
return {"run_id": "run-budget", "thread_id": kwargs["thread_id"]}
|
||||
|
||||
row = _once_task_row(task_id="task-budget", status="enabled")
|
||||
row.update({"schedule_type": "cron", "schedule_spec": {"cron": "* * * * *"}, "overlap_policy": "skip"})
|
||||
row.update({"schedule_type": "cron", "schedule_spec": {"cron": "* * * * *"}, "overlap_policy": "enqueue"})
|
||||
task_repo = DummyTaskRepo([row])
|
||||
# active_count equals max_concurrent_runs → budget is exhausted
|
||||
run_repo = DummyRunRepo(active_count=3)
|
||||
@ -1028,10 +1113,9 @@ async def test_manual_trigger_rejected_when_global_budget_exhausted():
|
||||
|
||||
result = await service.dispatch_task(row, now=datetime.now(UTC), trigger="manual")
|
||||
|
||||
assert result["outcome"] == "conflict"
|
||||
assert "limit" in result["error"]
|
||||
assert result["outcome"] == "queued"
|
||||
assert launched == []
|
||||
assert run_repo.created is None
|
||||
assert run_repo.created["status"] == "queued"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -1044,7 +1128,7 @@ async def test_manual_trigger_proceeds_when_global_budget_available():
|
||||
return {"run_id": "run-ok", "thread_id": kwargs["thread_id"]}
|
||||
|
||||
row = _once_task_row(task_id="task-ok", status="enabled")
|
||||
row.update({"schedule_type": "cron", "schedule_spec": {"cron": "* * * * *"}, "overlap_policy": "skip"})
|
||||
row.update({"schedule_type": "cron", "schedule_spec": {"cron": "* * * * *"}, "overlap_policy": "enqueue"})
|
||||
task_repo = DummyTaskRepo([row])
|
||||
# active_count is 2, max_concurrent_runs is 3 → one slot left
|
||||
run_repo = DummyRunRepo(active_count=2)
|
||||
|
||||
@ -2134,7 +2134,8 @@ agent_storage:
|
||||
# multi_instance: false # Opt into lease-aware recovery across Gateway instances
|
||||
# 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/reservations across multi-instance Pods
|
||||
# max_concurrent_runs: 3 # Global cap on launching/running scheduled runs across multi-instance Pods
|
||||
# queue_timeout_seconds: 3600 # Maximum durable queue wait before an occurrence fails
|
||||
# min_once_delay_seconds: 60 # Minimum future offset for one-time tasks at creation time
|
||||
# recursion_limit: 1000 # LangGraph super-step cap for scheduled runs (matches the web UI)
|
||||
scheduler:
|
||||
@ -2143,6 +2144,7 @@ scheduler:
|
||||
poll_interval_seconds: 5
|
||||
lease_seconds: 120
|
||||
max_concurrent_runs: 3
|
||||
queue_timeout_seconds: 3600
|
||||
min_once_delay_seconds: 60
|
||||
recursion_limit: 1000
|
||||
|
||||
|
||||
@ -291,11 +291,11 @@ The helper takes:
|
||||
|
||||
### Overlap semantics
|
||||
|
||||
MVP uses one fixed overlap rule:
|
||||
The original MVP used one fixed overlap rule:
|
||||
|
||||
- if the target thread already has a pending/running run, record the occurrence as `skipped`, update `next_run_at`, and do not launch another run.
|
||||
|
||||
This is intentionally narrower than exposing multiple overlap policies in the first PR.
|
||||
This historical rule was superseded by the durable `enqueue` behavior documented in the current README: a busy occurrence waits in `queued`, survives Gateway restarts, and fails only after the configured queue timeout.
|
||||
|
||||
### Misfire semantics
|
||||
|
||||
@ -522,7 +522,7 @@ The UI must not show a healthy-looking task row when the last scheduler attempt
|
||||
4. next-run computation across timezone boundaries and DST-sensitive cases
|
||||
5. one-time schedule status transitions
|
||||
6. due-task claim logic and lease expiry
|
||||
7. overlap skip behavior
|
||||
7. original overlap skip behavior (superseded by durable queue coverage)
|
||||
8. misfire `run_once` behavior
|
||||
|
||||
### Backend integration tests
|
||||
|
||||
@ -1,8 +1,10 @@
|
||||
"use client";
|
||||
|
||||
import { TriangleAlertIcon } from "lucide-react";
|
||||
import { useSearchParams } from "next/navigation";
|
||||
import { useEffect, useState } from "react";
|
||||
|
||||
import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
Dialog,
|
||||
@ -44,6 +46,22 @@ import { cn } from "@/lib/utils";
|
||||
|
||||
const NONE = "—";
|
||||
|
||||
function ReuseThreadNotice({
|
||||
title,
|
||||
description,
|
||||
}: {
|
||||
title: string;
|
||||
description: string;
|
||||
}) {
|
||||
return (
|
||||
<Alert className="border-amber-500/50 bg-amber-500/10">
|
||||
<TriangleAlertIcon className="text-amber-600 dark:text-amber-400" />
|
||||
<AlertTitle>{title}</AlertTitle>
|
||||
<AlertDescription>{description}</AlertDescription>
|
||||
</Alert>
|
||||
);
|
||||
}
|
||||
|
||||
function formatTimestamp(value: string | null, locale: string): string {
|
||||
if (!value) {
|
||||
return NONE;
|
||||
@ -236,11 +254,17 @@ export default function ScheduledTasksPage() {
|
||||
</Button>
|
||||
</div>
|
||||
{contextMode === "reuse_thread" && (
|
||||
<Input
|
||||
value={targetThreadId}
|
||||
onChange={(event) => setTargetThreadId(event.target.value)}
|
||||
placeholder={st.context.threadIdPlaceholder}
|
||||
/>
|
||||
<>
|
||||
<Input
|
||||
value={targetThreadId}
|
||||
onChange={(event) => setTargetThreadId(event.target.value)}
|
||||
placeholder={st.context.threadIdPlaceholder}
|
||||
/>
|
||||
<ReuseThreadNotice
|
||||
title={st.context.reuseNoticeTitle}
|
||||
description={st.context.reuseNoticeDescription}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
<Input
|
||||
value={title}
|
||||
@ -440,6 +464,12 @@ export default function ScheduledTasksPage() {
|
||||
? `${st.detail.thread}: ${selectedTask.thread_id ?? NONE}`
|
||||
: `${st.detail.lastThread}: ${selectedTask.last_thread_id ?? NONE}`}
|
||||
</div>
|
||||
{selectedTask.context_mode === "reuse_thread" && (
|
||||
<ReuseThreadNotice
|
||||
title={st.context.reuseNoticeTitle}
|
||||
description={st.context.reuseNoticeDescription}
|
||||
/>
|
||||
)}
|
||||
<div className="text-muted-foreground text-sm">
|
||||
{st.detail.schedule}:{" "}
|
||||
{scheduleTypeLabel(selectedTask.schedule_type)}
|
||||
|
||||
@ -391,6 +391,9 @@ export const enUS: Translations = {
|
||||
fresh: "Fresh thread",
|
||||
reuse: "Reuse thread",
|
||||
threadIdPlaceholder: "Thread ID",
|
||||
reuseNoticeTitle: "Uses this thread's conversation history",
|
||||
reuseNoticeDescription:
|
||||
"If this thread has an active run at the scheduled time, DeerFlow queues this occurrence and starts it when the thread is available. It fails if the configured queue wait limit is exceeded.",
|
||||
},
|
||||
filters: {
|
||||
allStatuses: "All statuses",
|
||||
@ -452,6 +455,7 @@ export const enUS: Translations = {
|
||||
runTrigger: { scheduled: "scheduled", manual: "manual" },
|
||||
runStatus: {
|
||||
queued: "Queued",
|
||||
launching: "Launching",
|
||||
running: "Running",
|
||||
success: "Success",
|
||||
failed: "Failed",
|
||||
|
||||
@ -310,6 +310,8 @@ export interface Translations {
|
||||
fresh: string;
|
||||
reuse: string;
|
||||
threadIdPlaceholder: string;
|
||||
reuseNoticeTitle: string;
|
||||
reuseNoticeDescription: string;
|
||||
};
|
||||
filters: {
|
||||
allStatuses: string;
|
||||
@ -370,6 +372,7 @@ export interface Translations {
|
||||
runTrigger: { scheduled: string; manual: string };
|
||||
runStatus: {
|
||||
queued: string;
|
||||
launching: string;
|
||||
running: string;
|
||||
success: string;
|
||||
failed: string;
|
||||
|
||||
@ -373,6 +373,9 @@ export const zhCN: Translations = {
|
||||
fresh: "新线程",
|
||||
reuse: "复用线程",
|
||||
threadIdPlaceholder: "线程 ID",
|
||||
reuseNoticeTitle: "使用该线程的历史对话",
|
||||
reuseNoticeDescription:
|
||||
"如果触发时该线程正在运行,DeerFlow 会将本次执行排队,并在线程空闲后启动;超过配置的最长等待时间后会标记为失败。",
|
||||
},
|
||||
filters: {
|
||||
allStatuses: "全部状态",
|
||||
@ -433,6 +436,7 @@ export const zhCN: Translations = {
|
||||
runTrigger: { scheduled: "定时", manual: "手动" },
|
||||
runStatus: {
|
||||
queued: "排队中",
|
||||
launching: "启动中",
|
||||
running: "运行中",
|
||||
success: "成功",
|
||||
failed: "失败",
|
||||
|
||||
@ -33,12 +33,14 @@ export type ScheduledTaskRun = {
|
||||
trigger: "scheduled" | "manual";
|
||||
status:
|
||||
| "queued"
|
||||
| "launching"
|
||||
| "running"
|
||||
| "success"
|
||||
| "failed"
|
||||
| "skipped"
|
||||
| "interrupted";
|
||||
error: string | null;
|
||||
attempt_count: number;
|
||||
started_at: string | null;
|
||||
finished_at: string | null;
|
||||
created_at: string;
|
||||
|
||||
@ -94,6 +94,61 @@ test("user can create a scheduled task from the page", async ({ page }) => {
|
||||
).toBeVisible();
|
||||
});
|
||||
|
||||
test("reuse-thread tasks explain their context and busy-thread queue behavior", async ({
|
||||
page,
|
||||
}) => {
|
||||
mockLangGraphAPI(page, {
|
||||
threads: [],
|
||||
scheduledTasks: [
|
||||
{
|
||||
id: "task-reuse",
|
||||
thread_id: "thread-1",
|
||||
context_mode: "reuse_thread",
|
||||
title: "Conversation summary",
|
||||
prompt: "Summarize this conversation",
|
||||
schedule_type: "cron",
|
||||
schedule_spec: { cron: "0 18 * * *" },
|
||||
timezone: "UTC",
|
||||
status: "enabled",
|
||||
next_run_at: "2026-07-02T18:00:00+00:00",
|
||||
last_run_at: null,
|
||||
last_run_id: null,
|
||||
last_error: null,
|
||||
run_count: 0,
|
||||
created_at: "2026-07-01T00:00:00+00:00",
|
||||
updated_at: "2026-07-01T00:00:00+00:00",
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
await page.goto("/workspace/scheduled-tasks");
|
||||
|
||||
const detailNotice = page
|
||||
.getByTestId("scheduled-task-detail")
|
||||
.getByRole("alert");
|
||||
await expect(detailNotice).toContainText(
|
||||
"Uses this thread's conversation history",
|
||||
);
|
||||
await expect(detailNotice).toContainText(
|
||||
"queues this occurrence and starts it when the thread is available",
|
||||
);
|
||||
|
||||
const createForm = page.getByTestId("scheduled-task-create-form");
|
||||
await expect(createForm.getByRole("alert")).toHaveCount(0);
|
||||
await createForm.getByRole("button", { name: "Reuse thread" }).click();
|
||||
|
||||
const createNotice = createForm.getByRole("alert");
|
||||
await expect(createNotice).toContainText(
|
||||
"Uses this thread's conversation history",
|
||||
);
|
||||
await expect(createNotice).toContainText(
|
||||
"It fails if the configured queue wait limit is exceeded",
|
||||
);
|
||||
|
||||
await createForm.getByRole("button", { name: "Fresh thread" }).click();
|
||||
await expect(createForm.getByRole("alert")).toHaveCount(0);
|
||||
});
|
||||
|
||||
test("user can pause a scheduled task from the detail pane", async ({
|
||||
page,
|
||||
}) => {
|
||||
|
||||
@ -256,8 +256,16 @@ export function mockLangGraphAPI(page: Page, options?: MockAPIOptions) {
|
||||
run_id: string | null;
|
||||
scheduled_for: string;
|
||||
trigger: "scheduled" | "manual";
|
||||
status: "queued" | "running" | "success" | "failed" | "skipped";
|
||||
status:
|
||||
| "queued"
|
||||
| "launching"
|
||||
| "running"
|
||||
| "success"
|
||||
| "failed"
|
||||
| "skipped"
|
||||
| "interrupted";
|
||||
error: string | null;
|
||||
attempt_count: number;
|
||||
started_at: string | null;
|
||||
finished_at: string | null;
|
||||
created_at: string;
|
||||
@ -539,6 +547,7 @@ export function mockLangGraphAPI(page: Page, options?: MockAPIOptions) {
|
||||
trigger: "manual",
|
||||
status: "success",
|
||||
error: null,
|
||||
attempt_count: 1,
|
||||
started_at: "2026-07-01T00:00:00+00:00",
|
||||
finished_at: "2026-07-01T00:00:00+00:00",
|
||||
created_at: "2026-07-01T00:00:00+00:00",
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user