mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-14 08:00:10 +00:00
cd0e74edaf
* fix(scheduler): reconcile stuck once tasks from committed run outcome Restart recovery (cancel_stuck_once_tasks and the multi-instance reconcile_stuck_once_tasks) blindly flipped every stuck once-task to 'cancelled'. When handle_run_completion crashed between its two transactions, a once-task whose run had already committed 'success' was permanently reported as cancelled. Both reconciliation paths now read the latest scheduled_task_runs row without a status filter and finalize the parent to match: success -> completed (last_error cleared), failed -> failed with the run's error, interrupted -> cancelled with the run's error when present, skipped -> cancelled (no work performed). Active occurrences (queued/launching/running) are left untouched — a concurrent completion or a later recovery pass will finalize them once the run reaches a terminal state. Tasks without a terminal run row keep the previous generic cancellation. Review follow-ups (willem-bd / Huixin615): - Extract _finalise_once_task_from_run() so both recovery paths share one outcome mapping (no more drift between single- and multi-instance paths). Returns bool (True = finalised, False = active/no-op) for explicit counter management at call sites. - Fix a no-op (`run_row.error or None` -> `run_row.error`) in the skipped branch. - Drop the unused `status` parameter from the test task helpers. - Use TERMINAL_RUN_STATUSES / ACTIVE_RUN_STATUSES constants (local copies to avoid circular import; kept in sync with scheduled_task_runs.sql). - [P1] Read the latest run AFTER acquiring the parent task row lock, not from a pre-lock batch snapshot. The latest-run lookup now runs per task under the lock with populate_existing so a concurrently committed status is read back fresh. - [P2] Race tests now use monkeypatch to actually enter the race window: _intercepted_fetch commits success in a separate session at the moment the per-task fetch fires, so a reverted pre-lock batch implementation fails the test, while the current post-lock implementation passes. - [P1] Do not finalize parent for active occurrences. A non-terminal scheduled occurrence means the run is still in progress — the parent must be left untouched until the completion path or a later recovery pass establishes a terminal outcome. - [P2] Add cancel_stuck_once_tasks to the single-instance poll loop so stuck once-tasks are not left permanently "running" when the startup sweep fails (mirrors multi-instance _reconcile_active_state behavior). - Fix stale docstrings in cancel_stuck_once_tasks and _fetch_latest_run. Adds regression tests for multiple historical runs (older success + newer skipped/active) on both paths, monkeypatch-based race tests that prove a concurrent completion committing success is reflected as completed, and active-run tests that verify the parent is left unchanged. Documents the behavior in AGENTS.md. Fixes #5034 * fix(scheduler): address review comments on completion-consistency fix - _fetch_latest_run: drop arbitrary id DESC tie-break; order by scheduled_for DESC (deterministic recency on schedule position) - _finalise_once_task_from_run: annotate bool return type - Centralize TERMINAL/ACTIVE_RUN_STATUSES in scheduled_tasks/model.py; stop duplicating them in scheduled_tasks/sql.py and scheduled_task_runs/sql.py (removes stale circular-import workaround) - cancel_stuck_once_tasks: run unconditionally in single-instance poll loop (remove try/except swallow) - tests: pin created_at/scheduled_for in _create_run so recency ordering is actually exercised; correct docstrings that described the active-occurrence branch as 'generic cancel' instead of 'left unchanged' * fix(scheduler): correct finalizer return annotation * fix: order scheduled task runs by creation time * fix(scheduler): stabilize latest run reconciliation ordering * fix(scheduler): order latest runs by creation time * test: update trace scheduler stub * fix(scheduler): clarify reconciliation diagnostics Signed-off-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com> * fix(scheduler): fail closed on startup recovery Keep single-instance parent reconciliation at startup so it cannot race manual admission. Propagate recovery failures through the Gateway lifespan before channel startup, preventing a half-started scheduler. Tests cover both recovery failure stages and a queued occurrence that survives startup before the ordinary poll drain launches it. Signed-off-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com> * fix(scheduler): order occurrences and fence stale parent writes Allocate per-task occurrence sequences under the parent lock and guard parent projection across launch, recovery, completion, and queue failure paths. Track launch accounting separately so stale occurrences are counted once without replacing newer results. Commit completion and accounting atomically, preserve legacy history, and cover migrations and reordered execution on SQLite and PostgreSQL. * fix(scheduler): tighten completion projection and launch fencing diagnostics Share the once-task outcome mapping between completion and both recovery paths, validate the terminal status before opening the completion transaction, leave cron parent status untouched on completion, log the fenced launch update when an occurrence does not belong to the launched run, and drop the README capability line. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * fix(scheduler): compare caller time only against unsequenced occurrences Among sequenced rows the parent-locked occurrence_seq is the only recency key. An unsequenced row can only be legacy history or an admission by a pre-upgrade Gateway writer, so recovery prefers it over the sequence winner only when its caller timestamp is later, which is the previous ordering for that pair. A rolling upgrade therefore degrades to the pre-sequence behaviour instead of ranking every pre-upgrade admission below every sequenced one. Document that boundary instead of requiring every Gateway writer to stop before the upgrade. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * fix(scheduler): gate once-task recovery on the same projection rule Recovery now finalises a once-task parent only from the occurrence that can_project() accepts: the highest sequenced occurrence whenever one exists, or the timestamp-latest row for a task whose history is entirely unsequenced. An unsequenced row admitted by a pre-upgrade writer can no longer cancel a parent whose sequenced occurrence is still live, nor stall finalisation of a parent whose sequenced occurrence already completed. Document that pre-upgrade instances project their own admissions during a rolling upgrade. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * fix(scheduler): defer once-task recovery while any occurrence is live uq_scheduled_task_run_active allows one non-terminal occurrence per task, so a live row is the task's newest admission whatever its caller clock and whether it carries a sequence. Both once-task recovery paths now probe for any active occurrence after the fresh latest-run read and leave the parent untouched while one exists; cancel_stuck_once_tasks also locks the parent row so admission cannot insert a queued occurrence between that probe and the commit. Once no occurrence is live, the sequence winner decides and a terminalised unsequenced row never overrides it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * test(persistence): follow the local head past canonical 0019 Main's forward-revision tests assumed 0019_thread_incarnations was the local chain head. With 0022_scheduled_occurrence_seq chained after it, seed the canonical-0019 shape explicitly, assert the real head where a database is upgraded, derive the 0020 rollback binary's revision set from the ancestors of its head, and step the PostgreSQL restart scenario back to canonical 0019 before the rollback binary restarts. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * docs(migrations): describe the chain through 0022_scheduled_occurrence_seq The rolling-forward section still ended the local chain at canonical 0019; it now names 0022_scheduled_occurrence_seq as the head and lists it among the revisions the 0020 rollback-floor binary does not know. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * test(scheduler): accept CI's sync Postgres URL in occurrence fixtures CI hands over TEST_POSTGRES_URI as postgresql://...?sslmode=disable. The occurrence, ordering and 0022 migration fixtures built async engines from it directly, so SQLAlchemy chose psycopg2, which is not installed. Normalize the scheme to postgresql+asyncpg and drop libpq-only query keys, matching the existing 0019 migration tests. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * docs(scheduler): keep the backend AGENTS.md chain within its budget The middlewares guidance chain was already above the hard limit on main, so any added byte in backend/AGENTS.md fails the agent guidance check. Leave backend/AGENTS.md identical to main and record the recovery projection rule in the 0022 migration entry, which already describes the occurrence fields. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Signed-off-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com> Co-authored-by: Totoro-qaq <279883115+Totoro-qaq@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
986 lines
35 KiB
Python
986 lines
35 KiB
Python
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_queued_once_task_survives_startup_and_is_drained_on_next_poll(tmp_path):
|
|
"""A newer queued occurrence survives recovery of an already-stuck parent.
|
|
|
|
The old success models a completion that committed before its parent update.
|
|
A real manual dispatch then leaves newer work queued after a transient
|
|
same-thread conflict. Startup must not let the old success finalize the
|
|
parent through the newer active row; the ordinary queue drain owns launch.
|
|
"""
|
|
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
|
try:
|
|
sf = get_session_factory()
|
|
assert sf is not None
|
|
task_repo = ScheduledTaskRepository(sf)
|
|
run_repo = ScheduledTaskRunRepository(sf)
|
|
now = datetime.now(UTC)
|
|
await task_repo.create(
|
|
task_id="task-queued-once",
|
|
user_id="user-1",
|
|
thread_id="thread-queued-once",
|
|
context_mode="reuse_thread",
|
|
assistant_id="lead_agent",
|
|
title="Queued once task",
|
|
prompt="Resume queued work",
|
|
schedule_type="once",
|
|
schedule_spec={"run_at": now.isoformat()},
|
|
timezone="UTC",
|
|
next_run_at=None,
|
|
)
|
|
await task_repo.update(
|
|
"task-queued-once",
|
|
user_id="user-1",
|
|
updates={"status": "running"},
|
|
)
|
|
await run_repo.create(
|
|
run_record_id="task-run-old-success",
|
|
task_id="task-queued-once",
|
|
thread_id="thread-queued-once",
|
|
scheduled_for=now - timedelta(minutes=1),
|
|
trigger="scheduled",
|
|
status="success",
|
|
)
|
|
task = await task_repo.get("task-queued-once", user_id="user-1")
|
|
assert task is not None
|
|
launched = []
|
|
|
|
async def launch_run(**kwargs):
|
|
launched.append(kwargs)
|
|
if len(launched) == 1:
|
|
raise ConflictError("Thread thread-queued-once already has an active run")
|
|
return {"run_id": "run-after-restart", "thread_id": kwargs["thread_id"]}
|
|
|
|
first_service = _make_service(task_repo, run_repo, launch_run)
|
|
queued = await first_service.dispatch_task(task, now=now, trigger="manual")
|
|
assert queued["outcome"] == "queued"
|
|
rows = await run_repo.list_by_task("task-queued-once")
|
|
assert [row["status"] for row in rows] == ["queued", "success"]
|
|
|
|
service = _make_service(task_repo, run_repo, launch_run)
|
|
|
|
async def parked_run_loop():
|
|
await service._stop.wait()
|
|
|
|
service._run_loop = parked_run_loop
|
|
await service.start()
|
|
try:
|
|
task = await task_repo.get_internal("task-queued-once")
|
|
rows = await run_repo.list_by_task("task-queued-once")
|
|
assert task is not None
|
|
assert task["status"] == "running"
|
|
assert rows[0]["status"] == "queued"
|
|
|
|
await service.run_once(now=now + timedelta(seconds=1))
|
|
|
|
task = await task_repo.get_internal("task-queued-once")
|
|
rows = await run_repo.list_by_task("task-queued-once")
|
|
assert task is not None
|
|
assert task["status"] == "running"
|
|
assert task["last_run_id"] == "run-after-restart"
|
|
assert rows[0]["status"] == "running"
|
|
assert rows[0]["run_id"] == "run-after-restart"
|
|
assert rows[1]["status"] == "success"
|
|
assert len(launched) == 2
|
|
finally:
|
|
await service.stop()
|
|
finally:
|
|
await close_engine()
|
|
|
|
|
|
async def test_only_one_worker_can_claim_a_queued_run(tmp_path):
|
|
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()
|