mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-15 15:50:35 +00:00
c263b887b5
* fix(scheduler): isolate queued run drain failures * fix(scheduler): isolate queue timeout rows * fix(scheduler): enforce queue timeout at launch claim Keep expired queued occurrences ineligible to start when the timeout sweep fails. Cover per-row and whole-scan cleanup failures with SQLite service tests that also verify healthy queued and due work still launches. * fix(scheduler): exclude expired rows from queue selection Filter rows past the queue deadline before owner ranking and the fair batch limit so failed timeout cleanup cannot starve healthy waiting work. Keep the atomic claim cutoff as the final launch guard. --------- Co-authored-by: AochenShen99 <sac1063467113@hotmail.com> Co-authored-by: Nan Gao <thunderggnn@gmail.com>
1523 lines
58 KiB
Python
1523 lines
58 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from uuid import UUID
|
|
|
|
import httpx
|
|
import pytest
|
|
import pytest_asyncio
|
|
from _router_auth_helpers import make_authed_test_app
|
|
from sqlalchemy import func, select, update
|
|
|
|
from app.gateway.auth.models import User
|
|
from app.gateway.auth_disabled import AUTH_SOURCE_SESSION
|
|
from app.gateway.routers import scheduled_tasks
|
|
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_task_runs.model import ScheduledTaskRunRow
|
|
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
|
|
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow
|
|
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_global_launch_budget_holds_when_distinct_rows_are_claimed_concurrently(tmp_path):
|
|
"""The global launch budget must survive claims racing on *distinct* rows.
|
|
|
|
``claim_queued_run`` counts executing rows and then promotes one row to
|
|
``launching``. Postgres serializes that pair with an advisory lock. SQLite
|
|
needs ``BEGIN IMMEDIATE`` for the same reason ``ThreadMetaRepository``
|
|
does: a deferred transaction does not reserve the writer until the UPDATE,
|
|
so every claimer reads the same stale count and overshoots
|
|
``max_concurrent_runs``.
|
|
|
|
Distinct rows are the load-bearing part. Two claims of the *same* row are
|
|
already safe via the ``status == "queued"`` CAS, which is what
|
|
``test_only_one_worker_can_claim_a_queued_run`` covers.
|
|
"""
|
|
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)
|
|
claimants = 8
|
|
for index in range(claimants):
|
|
await run_repo.create(
|
|
run_record_id=f"task-run-budget-{index}",
|
|
task_id=f"task-budget-{index}",
|
|
thread_id=f"thread-budget-{index}",
|
|
scheduled_for=now,
|
|
trigger="scheduled",
|
|
status="queued",
|
|
)
|
|
|
|
# Open every connection the claimers need up front. On a cold pool the
|
|
# per-connection PRAGMA setup staggers them enough to hide the race.
|
|
await asyncio.gather(*(run_repo.count_active_runs() for _ in range(claimants)))
|
|
|
|
# That warm-up is load-bearing, and it only works because the SQLite
|
|
# engine keeps pooled connections (init_engine_from_config builds it on
|
|
# SQLAlchemy's default AsyncAdaptedQueuePool, so `pool_size` of them
|
|
# survive this gather while the overflow is discarded). Under a
|
|
# non-pooling class such as NullPool every claimer would open its own
|
|
# connection, the PRAGMA setup would serialize them, and this test
|
|
# would pass against an unserialized claim instead of failing. Assert
|
|
# the reuse so that a pool change breaks this test loudly rather than
|
|
# quietly draining it of guard strength.
|
|
pool = sf.kw["bind"].sync_engine.pool
|
|
# A non-pooling class does not implement checkedin() at all, so treat a
|
|
# missing counter as "nothing was reused" and report it the same way.
|
|
pooled = pool.checkedin() if hasattr(pool, "checkedin") else 0
|
|
assert pooled >= 2, f"{type(pool).__name__} left {pooled} connections pooled after the warm-up; the claimers cannot overlap, so this test would pass against an unserialized claim"
|
|
|
|
claims = await asyncio.gather(
|
|
*(
|
|
run_repo.claim_queued_run(
|
|
f"task-run-budget-{index}",
|
|
lease_owner=f"worker-{index}",
|
|
now=now,
|
|
lease_seconds=120,
|
|
global_max_concurrent_runs=1,
|
|
)
|
|
for index in range(claimants)
|
|
)
|
|
)
|
|
|
|
assert sum(claim is not None for claim in claims) == 1
|
|
assert await run_repo.count_active_runs() == 1
|
|
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()
|
|
|
|
|
|
@pytest.mark.parametrize("failure_mode", ["row", "scan"])
|
|
async def test_queue_expiry_failure_does_not_launch_expired_work(
|
|
sqlite_repos,
|
|
monkeypatch,
|
|
caplog,
|
|
failure_mode,
|
|
):
|
|
"""A failed timeout sweep may delay cleanup, but never admits old work."""
|
|
sf, task_repo, run_repo = sqlite_repos
|
|
poll_at = datetime.now(UTC).replace(microsecond=0)
|
|
await _owner_task(
|
|
task_repo,
|
|
"task-expired",
|
|
user_id=OWNER_A,
|
|
next_run_at=poll_at - timedelta(minutes=5),
|
|
)
|
|
await _owner_task(
|
|
task_repo,
|
|
"task-healthy-queued",
|
|
user_id=OWNER_A,
|
|
next_run_at=poll_at + timedelta(hours=1),
|
|
)
|
|
await _owner_task(
|
|
task_repo,
|
|
"task-due",
|
|
user_id=OWNER_A,
|
|
next_run_at=poll_at - timedelta(minutes=1),
|
|
)
|
|
await _owned_occurrence(task_repo, run_repo, "expired", user_id=OWNER_A, task_id="task-expired")
|
|
await _owned_occurrence(
|
|
task_repo,
|
|
run_repo,
|
|
"healthy-queued",
|
|
user_id=OWNER_A,
|
|
task_id="task-healthy-queued",
|
|
)
|
|
async with sf() as session:
|
|
await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == "expired").values(created_at=poll_at - timedelta(hours=2)))
|
|
await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == "healthy-queued").values(created_at=poll_at - timedelta(minutes=1)))
|
|
await session.commit()
|
|
|
|
if failure_mode == "scan":
|
|
|
|
async def fail_expiry_scan(**_kwargs):
|
|
raise RuntimeError("queue expiry scan failed")
|
|
|
|
monkeypatch.setattr(run_repo, "expire_queued_runs", fail_expiry_scan)
|
|
else:
|
|
original_lock_task = run_repo._lock_task
|
|
failed = False
|
|
|
|
async def fail_expired_row_once(session, task_id):
|
|
nonlocal failed
|
|
if task_id == "task-expired" and not failed:
|
|
failed = True
|
|
raise RuntimeError("queue expiry row failed")
|
|
return await original_lock_task(session, task_id)
|
|
|
|
monkeypatch.setattr(run_repo, "_lock_task", fail_expired_row_once)
|
|
|
|
launches: list[dict] = []
|
|
service = _capped_service(task_repo, run_repo, launches, max_concurrent_runs=3, queue_timeout_seconds=3600)
|
|
|
|
await service.run_once(now=poll_at)
|
|
|
|
launched_tasks = {item["metadata"]["scheduled_task_id"] for item in launches}
|
|
assert launched_tasks == {"task-healthy-queued", "task-due"}
|
|
expired = (await run_repo.list_by_task("task-expired"))[0]
|
|
assert (expired["status"], expired["error"], expired["attempt_count"]) == (
|
|
"queued",
|
|
None,
|
|
0,
|
|
)
|
|
assert "Failed to expire queued scheduled" in caplog.text
|
|
|
|
|
|
async def test_expired_queue_rows_do_not_consume_capped_candidates(
|
|
sqlite_repos,
|
|
monkeypatch,
|
|
):
|
|
"""Timeout failures must not let expired rows crowd healthy work out of the batch."""
|
|
sf, task_repo, run_repo = sqlite_repos
|
|
poll_at = datetime.now(UTC).replace(microsecond=0)
|
|
for task_id, row_id, created_at in (
|
|
("task-expired-1", "expired-1", poll_at - timedelta(hours=2)),
|
|
("task-expired-2", "expired-2", poll_at - timedelta(hours=2)),
|
|
("task-healthy", "healthy-queued", poll_at - timedelta(minutes=1)),
|
|
):
|
|
await _owner_task(
|
|
task_repo,
|
|
task_id,
|
|
user_id=OWNER_A,
|
|
next_run_at=poll_at + timedelta(hours=1),
|
|
)
|
|
await _owned_occurrence(
|
|
task_repo,
|
|
run_repo,
|
|
row_id,
|
|
user_id=OWNER_A,
|
|
task_id=task_id,
|
|
)
|
|
async with sf() as session:
|
|
await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == row_id).values(created_at=created_at))
|
|
await session.commit()
|
|
|
|
async def fail_expiry_scan(**_kwargs):
|
|
raise RuntimeError("queue expiry scan failed")
|
|
|
|
monkeypatch.setattr(run_repo, "expire_queued_runs", fail_expiry_scan)
|
|
|
|
launches: list[dict] = []
|
|
service = _capped_service(
|
|
task_repo,
|
|
run_repo,
|
|
launches,
|
|
max_concurrent_runs=3,
|
|
max_concurrent_runs_per_user=2,
|
|
queue_timeout_seconds=3600,
|
|
)
|
|
|
|
for offset in range(3):
|
|
await service.run_once(now=poll_at + timedelta(seconds=offset))
|
|
|
|
assert [item["metadata"]["scheduled_task_id"] for item in launches] == ["task-healthy"]
|
|
for row_id in ("expired-1", "expired-2"):
|
|
row = (await run_repo.list_by_task(f"task-{row_id}"))[0]
|
|
assert (row["status"], row["attempt_count"]) == ("queued", 0)
|
|
healthy = (await run_repo.list_by_task("task-healthy"))[0]
|
|
assert healthy["status"] == "running"
|
|
|
|
|
|
async def test_expired_same_thread_head_does_not_hide_unexpired_candidate(sqlite_repos):
|
|
"""An expired FIFO head cannot block a newer row once it is due for timeout."""
|
|
sf, task_repo, run_repo = sqlite_repos
|
|
poll_at = datetime.now(UTC).replace(microsecond=0)
|
|
await _owned_occurrence(
|
|
task_repo,
|
|
run_repo,
|
|
"expired-head",
|
|
user_id=OWNER_A,
|
|
thread_id="shared-thread",
|
|
)
|
|
await _owned_occurrence(
|
|
task_repo,
|
|
run_repo,
|
|
"healthy-next",
|
|
user_id=OWNER_A,
|
|
thread_id="shared-thread",
|
|
)
|
|
async with sf() as session:
|
|
await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == "expired-head").values(created_at=poll_at - timedelta(hours=2)))
|
|
await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.id == "healthy-next").values(created_at=poll_at - timedelta(minutes=1)))
|
|
await session.commit()
|
|
|
|
batch = await run_repo.list_queued_runs(
|
|
limit=16,
|
|
created_after=poll_at - timedelta(hours=1),
|
|
)
|
|
|
|
assert [row["id"] for row in batch] == ["healthy-next"]
|
|
|
|
|
|
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()
|
|
|
|
|
|
# --- PR2 M1: per-owner run cap and fair queue selection ------------------------------------
|
|
|
|
_QUEUE_CONTRACT = json.loads((Path(__file__).resolve().parents[2] / "contracts" / "scheduled_goal_notes_contract.json").read_text(encoding="utf-8"))
|
|
_OWNER_A = User(id=UUID("aaaaaaaa-0000-4000-8000-000000000001"), email="owner-a@example.com", password_hash="unused", system_role="user")
|
|
OWNER_A = str(_OWNER_A.id)
|
|
OWNER_B = "bbbbbbbb-0000-4000-8000-000000000002"
|
|
|
|
|
|
async def _owner_task(task_repo: ScheduledTaskRepository, task_id: str, *, user_id: str, next_run_at: datetime | None = None) -> dict:
|
|
"""A page-created per-minute task with a fresh thread for every run."""
|
|
return await task_repo.create(
|
|
task_id=task_id,
|
|
user_id=user_id,
|
|
thread_id=None,
|
|
context_mode="fresh_thread_per_run",
|
|
assistant_id="lead_agent",
|
|
title=task_id,
|
|
prompt="Check the list",
|
|
schedule_type="cron",
|
|
schedule_spec={"cron": "* * * * *"},
|
|
timezone="UTC",
|
|
next_run_at=next_run_at,
|
|
)
|
|
|
|
|
|
async def _owned_occurrence(
|
|
task_repo: ScheduledTaskRepository,
|
|
run_repo: ScheduledTaskRunRepository,
|
|
row_id: str,
|
|
*,
|
|
user_id: str,
|
|
status: str = "queued",
|
|
thread_id: str | None = None,
|
|
task_id: str | None = None,
|
|
) -> dict:
|
|
"""One occurrence on its own task (one active occurrence per task)."""
|
|
task_id = task_id or f"task-{row_id}"
|
|
if await task_repo.get_internal(task_id) is None:
|
|
await _owner_task(task_repo, task_id, user_id=user_id)
|
|
return await run_repo.create(
|
|
run_record_id=row_id,
|
|
task_id=task_id,
|
|
thread_id=thread_id or f"thread-{row_id}",
|
|
scheduled_for=datetime.now(UTC),
|
|
trigger="scheduled",
|
|
status=status,
|
|
)
|
|
|
|
|
|
async def _executing_by_owner(sf, owner: str) -> int:
|
|
async with sf() as session:
|
|
return int(
|
|
await session.scalar(
|
|
select(func.count())
|
|
.select_from(ScheduledTaskRunRow)
|
|
.join(ScheduledTaskRow, ScheduledTaskRow.id == ScheduledTaskRunRow.task_id)
|
|
.where(ScheduledTaskRunRow.status.in_(("launching", "running")), ScheduledTaskRow.user_id == owner)
|
|
)
|
|
or 0
|
|
)
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def sqlite_repos(tmp_path):
|
|
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
|
sf = get_session_factory()
|
|
assert sf is not None
|
|
try:
|
|
yield sf, ScheduledTaskRepository(sf), ScheduledTaskRunRepository(sf)
|
|
finally:
|
|
await close_engine()
|
|
|
|
|
|
def _capped_service(task_repo, run_repo, launches: list[dict], **limits) -> ScheduledTaskService:
|
|
"""The real service with the given caps; every launch is recorded in ``launches``."""
|
|
|
|
async def launch_run(**kwargs):
|
|
launches.append(kwargs)
|
|
return {"run_id": f"run-{len(launches)}", "thread_id": kwargs["thread_id"]}
|
|
|
|
return ScheduledTaskService(task_repo=task_repo, task_run_repo=run_repo, launch_run=launch_run, poll_interval_seconds=5, lease_seconds=120, **limits)
|
|
|
|
|
|
async def test_fair_batch_includes_other_owner_when_one_owner_is_at_cap(sqlite_repos):
|
|
_sf, task_repo, run_repo = sqlite_repos
|
|
for index in range(2):
|
|
await _owned_occurrence(task_repo, run_repo, f"a-running-{index}", user_id=OWNER_A, status="running")
|
|
for index in range(50):
|
|
await _owned_occurrence(task_repo, run_repo, f"a-queued-{index:02d}", user_id=OWNER_A)
|
|
await _owned_occurrence(task_repo, run_repo, "b-queued", user_id=OWNER_B)
|
|
|
|
batch = await run_repo.list_queued_runs(limit=16, per_user_max_concurrent_runs=2)
|
|
|
|
# A is at its cap, so its 50 waiting rows drop out before the limit.
|
|
assert [row["id"] for row in batch] == ["b-queued"]
|
|
|
|
|
|
async def test_fair_batch_interleaves_owners_without_cap(sqlite_repos):
|
|
_sf, task_repo, run_repo = sqlite_repos
|
|
for index in range(50):
|
|
await _owned_occurrence(task_repo, run_repo, f"a-queued-{index:02d}", user_id=OWNER_A)
|
|
await _owned_occurrence(task_repo, run_repo, "b-queued", user_id=OWNER_B)
|
|
|
|
batch = [row["id"] for row in await run_repo.list_queued_runs(limit=16)]
|
|
|
|
# Without a cap the batch is still fair: B's only row, created last,
|
|
# comes right after A's head instead of after 50 older A rows.
|
|
assert len(batch) == 16
|
|
assert batch.index("b-queued") <= 1
|
|
assert batch[0] == "a-queued-00"
|
|
# Within an owner the order is unchanged (attempt_count, created_at, id).
|
|
assert [row_id for row_id in batch if row_id.startswith("a-")] == [f"a-queued-{index:02d}" for index in range(15)]
|
|
|
|
|
|
async def test_owner_rank_respects_remaining_capacity(sqlite_repos):
|
|
_sf, task_repo, run_repo = sqlite_repos
|
|
await _owned_occurrence(task_repo, run_repo, "a-running", user_id=OWNER_A, status="running")
|
|
for index in range(5):
|
|
await _owned_occurrence(task_repo, run_repo, f"a-queued-{index}", user_id=OWNER_A)
|
|
await _owned_occurrence(task_repo, run_repo, "b-queued", user_id=OWNER_B)
|
|
|
|
batch = [row["id"] for row in await run_repo.list_queued_runs(limit=16, per_user_max_concurrent_runs=3)]
|
|
|
|
assert [row_id for row_id in batch if row_id.startswith("a-")] == ["a-queued-0", "a-queued-1"]
|
|
assert "b-queued" in batch
|
|
|
|
|
|
@pytest.mark.parametrize("cap", [0, 2])
|
|
async def test_fair_batch_keeps_same_thread_fifo(sqlite_repos, cap):
|
|
_sf, task_repo, run_repo = sqlite_repos
|
|
await _owned_occurrence(task_repo, run_repo, "a-older", user_id=OWNER_A, thread_id="reuse-thread")
|
|
await _owned_occurrence(task_repo, run_repo, "a-newer", user_id=OWNER_A, thread_id="reuse-thread")
|
|
|
|
batch = [row["id"] for row in await run_repo.list_queued_runs(limit=16, per_user_max_concurrent_runs=cap)]
|
|
|
|
assert batch == ["a-older"]
|
|
|
|
|
|
async def test_orphan_queued_row_is_still_drained(sqlite_repos):
|
|
_sf, task_repo, run_repo = sqlite_repos
|
|
await _owned_occurrence(task_repo, run_repo, "a-running", user_id=OWNER_A, status="running")
|
|
# The parent task is gone (deleted while the row waited): no owner.
|
|
await run_repo.create(run_record_id="orphan", task_id="task-deleted", thread_id="thread-orphan", scheduled_for=datetime.now(UTC), trigger="scheduled", status="queued")
|
|
|
|
assert [row["id"] for row in await run_repo.list_queued_runs(limit=16, per_user_max_concurrent_runs=1)] == ["orphan"]
|
|
|
|
launches: list[dict] = []
|
|
service = _capped_service(task_repo, run_repo, launches, max_concurrent_runs=3, max_concurrent_runs_per_user=1)
|
|
await service.run_once(now=datetime.now(UTC))
|
|
|
|
row = (await run_repo.list_by_task("task-deleted"))[0]
|
|
assert (row["status"], row["error"]) == ("interrupted", _QUEUE_CONTRACT["host_run_errors"]["deleted_while_queued"])
|
|
assert launches == []
|
|
|
|
|
|
async def test_claim_rejects_owner_at_cap_but_admits_other_owner(sqlite_repos):
|
|
_sf, task_repo, run_repo = sqlite_repos
|
|
now = datetime.now(UTC)
|
|
await _owned_occurrence(task_repo, run_repo, "a-running", user_id=OWNER_A, status="running")
|
|
await _owned_occurrence(task_repo, run_repo, "a-queued", user_id=OWNER_A)
|
|
await _owned_occurrence(task_repo, run_repo, "b-queued", user_id=OWNER_B)
|
|
claim = {"now": now, "lease_seconds": 120, "global_max_concurrent_runs": 3, "per_user_max_concurrent_runs": 1}
|
|
|
|
assert await run_repo.claim_queued_run("a-queued", lease_owner="worker", **claim) is None
|
|
assert (await run_repo.list_by_task("task-a-queued"))[0]["status"] == "queued"
|
|
assert (await run_repo.list_by_task("task-a-queued"))[0]["attempt_count"] == 0
|
|
|
|
claimed = await run_repo.claim_queued_run("b-queued", lease_owner="worker", **claim)
|
|
assert claimed is not None and claimed["status"] == "launching"
|
|
|
|
# Cap 0 keeps today's behavior: only the global budget applies.
|
|
assert await run_repo.claim_queued_run("a-queued", lease_owner="worker", **{**claim, "per_user_max_concurrent_runs": 0}) is not None
|
|
|
|
|
|
async def test_effective_owner_cap_is_bounded_by_global_cap(sqlite_repos):
|
|
sf, task_repo, run_repo = sqlite_repos
|
|
for index in range(3):
|
|
await _owned_occurrence(task_repo, run_repo, f"a-queued-{index}", user_id=OWNER_A)
|
|
launches: list[dict] = []
|
|
service = _capped_service(task_repo, run_repo, launches, max_concurrent_runs=1, max_concurrent_runs_per_user=2)
|
|
|
|
assert service._max_concurrent_runs_per_user == 1
|
|
await service.run_once(now=datetime.now(UTC))
|
|
|
|
assert len(launches) == 1
|
|
assert await run_repo.count_active_runs() == 1
|
|
assert await _executing_by_owner(sf, OWNER_A) == 1
|
|
# The drain asks only for as many A rows as A may still start.
|
|
assert await run_repo.list_queued_runs(limit=16, per_user_max_concurrent_runs=service._max_concurrent_runs_per_user) == []
|
|
|
|
|
|
async def test_capped_owner_queue_expiry_writes_contract_queue_timeout_note(sqlite_repos):
|
|
sf, task_repo, run_repo = sqlite_repos
|
|
await _owned_occurrence(task_repo, run_repo, "a-running", user_id=OWNER_A, status="running")
|
|
await _owned_occurrence(task_repo, run_repo, "a-waiting", user_id=OWNER_A)
|
|
launches: list[dict] = []
|
|
service = _capped_service(task_repo, run_repo, launches, max_concurrent_runs=3, max_concurrent_runs_per_user=1, queue_timeout_seconds=60)
|
|
|
|
# While A holds its only slot, the waiting row is not launched.
|
|
await service.run_once(now=datetime.now(UTC))
|
|
assert launches == []
|
|
assert (await run_repo.list_by_task("task-a-waiting"))[0]["status"] == "queued"
|
|
|
|
await service.run_once(now=datetime.now(UTC) + timedelta(seconds=61))
|
|
|
|
expired = (await run_repo.list_by_task("task-a-waiting"))[0]
|
|
assert (expired["status"], expired["run_id"]) == ("failed", None)
|
|
assert expired["error"] == _QUEUE_CONTRACT["host_run_errors"]["queue_timeout"]
|
|
assert launches == []
|
|
|
|
# The runs route hands that note to the page unchanged, so PR1's
|
|
# runErrors.queueTimeout mapping ("waited too long") applies.
|
|
app = make_authed_test_app(user_factory=lambda: _OWNER_A, bind_current_user=True)
|
|
|
|
@app.middleware("http")
|
|
async def mark_session_source(request, call_next):
|
|
request.state.auth_source = AUTH_SOURCE_SESSION
|
|
return await call_next(request)
|
|
|
|
app.state.scheduled_task_repo = task_repo
|
|
app.state.scheduled_task_run_repo = run_repo
|
|
app.include_router(scheduled_tasks.router)
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://queue.test") as client:
|
|
response = await client.get("/api/scheduled-tasks/task-a-waiting/runs")
|
|
assert response.status_code == 200, response.text
|
|
assert [(row["status"], row["run_id"], row["error"]) for row in response.json()] == [("failed", None, _QUEUE_CONTRACT["host_run_errors"]["queue_timeout"])]
|
|
|
|
|
|
async def test_service_drain_launches_other_owner_first(sqlite_repos):
|
|
"""Two poll cycles: A's backlog never takes B's slot.
|
|
|
|
``claim_due_tasks`` admits ``max_concurrent_runs`` tasks per poll ordered
|
|
by ``next_run_at, id``, so A's five tasks (due first) fill cycle 1. Each
|
|
admitted occurrence is offered a launch at once; the per-owner cap lets
|
|
one A row start and leaves the other two queued. In cycle 2 the drain
|
|
must not start any A row, and B's occurrence (due one second later)
|
|
launches although four older A rows are queued ahead of it.
|
|
"""
|
|
sf, task_repo, run_repo = sqlite_repos
|
|
t0 = datetime.now(UTC).replace(microsecond=0) - timedelta(minutes=5)
|
|
for index in range(5):
|
|
await _owner_task(task_repo, f"task-a-{index}", user_id=OWNER_A, next_run_at=t0)
|
|
await _owner_task(task_repo, "task-b", user_id=OWNER_B, next_run_at=t0 + timedelta(seconds=1))
|
|
launches: list[dict] = []
|
|
service = _capped_service(task_repo, run_repo, launches, max_concurrent_runs=3, max_concurrent_runs_per_user=1)
|
|
|
|
b_cycle = None
|
|
for cycle in range(1, 5):
|
|
await service.run_once(now=t0 + timedelta(seconds=cycle))
|
|
assert await _executing_by_owner(sf, OWNER_A) <= 1
|
|
b_rows = await run_repo.list_by_task("task-b")
|
|
if b_rows:
|
|
b_cycle = cycle
|
|
break
|
|
|
|
assert b_cycle == 2
|
|
assert b_rows[0]["status"] in {"launching", "running"}
|
|
assert await _executing_by_owner(sf, OWNER_A) == 1
|
|
queued_a = [row for index in range(5) for row in await run_repo.list_by_task(f"task-a-{index}") if row["status"] == "queued"]
|
|
assert len(queued_a) == 4
|
|
|
|
|
|
async def test_service_drain_launches_other_owner_from_a_seeded_queue(sqlite_repos):
|
|
"""A seeded backlog larger than the drain batch cannot hide B's row."""
|
|
sf, task_repo, run_repo = sqlite_repos
|
|
for index in range(20):
|
|
await _owned_occurrence(task_repo, run_repo, f"a-queued-{index:02d}", user_id=OWNER_A)
|
|
await _owned_occurrence(task_repo, run_repo, "b-queued", user_id=OWNER_B)
|
|
launches: list[dict] = []
|
|
service = _capped_service(task_repo, run_repo, launches, max_concurrent_runs=3, max_concurrent_runs_per_user=1)
|
|
|
|
await service.run_once(now=datetime.now(UTC))
|
|
|
|
# The batch limit is max(16, 3 * 4) = 16 < 20 A rows: an unfair batch
|
|
# would hold only A rows. The fair one starts A's head and B's row.
|
|
assert (await run_repo.list_by_task("task-b-queued"))[0]["status"] == "running"
|
|
assert (await run_repo.list_by_task("task-a-queued-00"))[0]["status"] == "running"
|
|
assert await _executing_by_owner(sf, OWNER_A) == 1
|
|
assert len(launches) == 2
|
|
|
|
|
|
@pytest.mark.parametrize("dialect_name", ["postgresql", "sqlite"])
|
|
async def test_fair_queue_statement_compiles_for_both_dialects(dialect_name):
|
|
"""The fair drain is plain SQLAlchemy Core (no raw text) on both dialects."""
|
|
from sqlalchemy.dialects import postgresql, sqlite
|
|
|
|
dialect = {"postgresql": postgresql.dialect(), "sqlite": sqlite.dialect()}[dialect_name]
|
|
sql = str(ScheduledTaskRunRepository._fair_queue_statement(limit=16, per_user_max_concurrent_runs=2).compile(dialect=dialect)).lower()
|
|
|
|
assert "row_number() over (partition by" in sql
|
|
assert "group by" in sql and "not (exists" in sql
|
|
assert "limit" in sql
|