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 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() 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