deer-flow/backend/tests/test_scheduled_task_queue.py

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