deer-flow/backend/tests/test_scheduled_task_postgres.py

from __future__ import annotations

import asyncio
import os
import uuid
from contextlib import asynccontextmanager
from datetime import UTC, datetime, timedelta
from unittest.mock import patch
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit

import pytest
import pytest_asyncio
from sqlalchemy import select, text, update
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from deerflow.config.database_config import DatabaseConfig
from deerflow.persistence.base import Base
from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config
from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
from deerflow.persistence.run import RunRepository
from deerflow.persistence.run.model import RunChangeClockRow, RunRow
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.events.store.db import DbRunEventStore
from deerflow.runtime.events.terminal import RecoveryTerminalPublisher

POSTGRES_URL = os.environ.get("TEST_POSTGRES_URI")

pytestmark = pytest.mark.skipif(
    not POSTGRES_URL,
    reason="requires TEST_POSTGRES_URI (real Postgres for scheduler concurrency)",
)


def _postgres_url(url: str) -> str:
    parts = urlsplit(url)
    query = urlencode([(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key not in {"sslmode", "channel_binding"}])
    return urlunsplit(parts._replace(query=query))


@pytest_asyncio.fixture()
async def postgres_repositories():
    assert POSTGRES_URL is not None
    schema = f"scheduler_{uuid.uuid4().hex}"
    await init_engine_from_config(
        DatabaseConfig(
            backend="postgres",
            postgres_url=_postgres_url(POSTGRES_URL),
            postgres_schema=schema,
        )
    )
    sf = get_session_factory()
    assert sf is not None
    run_repo = RunRepository(sf)
    task_repo = ScheduledTaskRepository(sf, run_repository=run_repo)
    task_run_repo = ScheduledTaskRunRepository(sf, run_repository=run_repo)
    try:
        yield task_repo, task_run_repo, run_repo
    finally:
        engine = get_engine()
        async with engine.begin() as connection:
            await connection.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
        await close_engine()


async def _create_cron_task(task_repo: ScheduledTaskRepository, task_id: str, *, next_run_at: datetime | None) -> None:
    await task_repo.create(
        task_id=task_id,
        user_id="user-1",
        thread_id=None,
        context_mode="fresh_thread_per_run",
        assistant_id="lead_agent",
        title=task_id,
        prompt="p",
        schedule_type="cron",
        schedule_spec={"cron": "* * * * *"},
        timezone="UTC",
        next_run_at=next_run_at,
    )


@pytest.mark.asyncio
async def test_postgres_global_budget_serializes_cross_pod_claims(postgres_repositories):
    task_repo, task_run_repo, _run_repo = postgres_repositories
    now = datetime.now(UTC)
    for suffix in ("a", "b"):
        await _create_cron_task(task_repo, f"task-{suffix}", next_run_at=now)
        await task_run_repo.create(
            run_record_id=f"task-run-{suffix}",
            task_id=f"task-{suffix}",
            thread_id=f"thread-{suffix}",
            scheduled_for=now,
            trigger="scheduled",
            status="queued",
        )

    claims = await asyncio.gather(
        task_run_repo.claim_queued_run(
            "task-run-a",
            now=now,
            lease_owner="pod-a",
            lease_seconds=60,
            global_max_concurrent_runs=1,
        ),
        task_run_repo.claim_queued_run(
            "task-run-b",
            now=now,
            lease_owner="pod-b",
            lease_seconds=60,
            global_max_concurrent_runs=1,
        ),
    )

    assert sum(claimed is not None for claimed in claims) == 1
    statuses = {(await task_run_repo.list_by_task(task_id))[0]["status"] for task_id in ("task-a", "task-b")}
    assert statuses == {"queued", "launching"}


@pytest.mark.asyncio
async def test_postgres_reconciliation_uses_metadata_and_atomically_claims_expired_run(postgres_repositories):
    task_repo, task_run_repo, run_repo = postgres_repositories
    now = datetime.now(UTC)
    for suffix in ("live", "expired"):
        await _create_cron_task(task_repo, f"task-{suffix}", next_run_at=None)
        await task_run_repo.create(
            run_record_id=f"task-run-{suffix}",
            task_id=f"task-{suffix}",
            thread_id=f"thread-{suffix}",
            scheduled_for=now,
            trigger="scheduled",
            status="running" if suffix == "expired" else "queued",
        )

    await task_run_repo.update_status("task-run-expired", status="running", run_id="run-expired")
    await run_repo.put(
        "run-live",
        thread_id="thread-live",
        user_id="user-1",
        status="running",
        metadata={
            "scheduled_task_id": "task-live",
            "scheduled_task_run_id": "task-run-live",
        },
        owner_worker_id="pod-a",
        lease_expires_at=(now + timedelta(seconds=60)).isoformat(),
    )
    await run_repo.put(
        "run-expired",
        thread_id="thread-expired",
        user_id="user-1",
        status="running",
        metadata={
            "scheduled_task_id": "task-expired",
            "scheduled_task_run_id": "task-run-expired",
        },
        owner_worker_id="pod-dead",
        lease_expires_at=(now - timedelta(seconds=60)).isoformat(),
    )

    assert await task_run_repo.reconcile_active_runs(error="lease expired", now=now) == 1
    assert (await task_run_repo.list_by_task("task-live"))[0]["status"] == "queued"
    assert (await task_run_repo.list_by_task("task-expired"))[0]["status"] == "interrupted"
    recovered = await run_repo.get("run-expired", user_id=None)
    assert recovered is not None
    assert recovered["status"] == "error"
    assert recovered["stop_reason"] == "scheduled_task_orphan_recovered"


@pytest.mark.asyncio
async def test_postgres_recovery_takeover_has_one_publisher():
    async with _contender_engines(2) as factories:
        seed_repo = RunRepository(factories[0])
        event_store = DbRunEventStore(factories[0])
        now = datetime.now(UTC)
        run_id, thread_id = "run-recovery-postgres", "thread-recovery-postgres"
        await seed_repo.put(
            run_id,
            thread_id=thread_id,
            user_id="user-1",
            status="running",
            owner_worker_id="dead-owner",
            lease_expires_at=(now - timedelta(seconds=60)).isoformat(),
        )

        both_ready = asyncio.Event()
        start_claiming = asyncio.Event()
        ready_count = 0
        ready_lock = asyncio.Lock()

        async def recover_once(factory):
            nonlocal ready_count
            repo = RunRepository(factory)
            publisher = RecoveryTerminalPublisher(DbRunEventStore(factory), repo)
            async with ready_lock:
                ready_count += 1
                if ready_count == 2:
                    both_ready.set()
            await both_ready.wait()
            await start_claiming.wait()
            claimed = await repo.claim_for_takeover(
                run_id,
                grace_seconds=10,
                error="orphan recovered",
                stop_reason="orphan_recovered",
            )
            target = publisher.authorize_takeover(
                claimed=claimed,
                thread_id=thread_id,
                run_id=run_id,
                user_id="user-1",
                operation_kind="run",
            )
            if target is not None:
                await publisher.publish(target)
            return target is not None

        tasks = [asyncio.create_task(recover_once(factory)) for factory in factories]
        try:
            await asyncio.wait_for(both_ready.wait(), timeout=5)
            start_claiming.set()
            outcomes = await asyncio.wait_for(asyncio.gather(*tasks), timeout=10)
        finally:
            start_claiming.set()
            await asyncio.gather(*tasks, return_exceptions=True)

        assert sum(outcomes) == 1
        terminal = await event_store.list_events(thread_id, run_id, event_types=["run.terminal"], user_id=None)
        assert len(terminal) == 1
        assert terminal[0]["content"] == {"status": "error"}
        assert terminal[0]["metadata"]["recovered"] is True


@pytest.mark.asyncio
@pytest.mark.parametrize("recovery_method", ["cancel_stuck_once_tasks", "reconcile_stuck_once_tasks"])
@pytest.mark.parametrize("newer_status", ["skipped", "queued", "launching", "running"])
@pytest.mark.parametrize("older_clock_ahead_seconds", [30, 0], ids=["reversed-timestamps", "equal-timestamps"])
async def test_postgres_once_recovery_uses_occurrence_order_despite_clock_skew(postgres_repositories, recovery_method, newer_status, older_clock_ahead_seconds):
    task_repo, task_run_repo, _run_repo = postgres_repositories
    now = datetime(2026, 7, 15, 12, 0, tzinfo=UTC)
    await task_repo.create(
        task_id="task-once",
        user_id="user-1",
        thread_id=None,
        context_mode="fresh_thread_per_run",
        assistant_id=None,
        title="Once task",
        prompt="p",
        schedule_type="once",
        schedule_spec={"run_at": now.isoformat()},
        timezone="UTC",
        next_run_at=None,
    )
    await task_repo.update_after_launch(
        "task-once",
        status="running",
        next_run_at=None,
        last_run_at=now,
        last_run_id=None,
        last_thread_id=None,
        last_error=None,
        increment_run_count=False,
    )
    # The older worker's clock and descending UUID order both favor its success.
    with patch("deerflow.persistence.scheduled_task_runs.sql.datetime") as clock:
        clock.now.return_value = now + timedelta(seconds=older_clock_ahead_seconds)
        older = await task_run_repo.create(
            run_record_id="ffffffff-ffff-4fff-8fff-ffffffffffff",
            task_id="task-once",
            thread_id="thread-old",
            scheduled_for=clock.now.return_value,
            trigger="manual",
            status="success",
        )
        clock.now.return_value = now
        newer = await task_run_repo.create(
            run_record_id="00000000-0000-4000-8000-000000000000",
            task_id="task-once",
            thread_id="thread-new",
            scheduled_for=clock.now.return_value,
            trigger="manual",
            status=newer_status,
        )
    assert older["created_at"] >= newer["created_at"]
    assert older["scheduled_for"] >= newer["scheduled_for"]
    kwargs = {"error": "interrupted: recovery"}
    if recovery_method == "reconcile_stuck_once_tasks":
        kwargs["now"] = now + timedelta(minutes=1)
    count = await getattr(task_repo, recovery_method)(**kwargs)

    task = await task_repo.get_internal("task-once")
    assert task is not None
    assert task["status"] == ("cancelled" if newer_status == "skipped" else "running")
    assert count == (1 if newer_status == "skipped" else 0)
    assert task["last_error"] is None
    async with task_run_repo._sf() as session:
        older_row = await session.get(ScheduledTaskRunRow, older["id"])
        newer_row = await session.get(ScheduledTaskRunRow, newer["id"])
        assert newer_row.occurrence_seq > older_row.occurrence_seq


@asynccontextmanager
async def _contender_engines(count: int):
    """``count`` independent engines on one fresh schema (one pool per contender)."""
    assert POSTGRES_URL is not None
    parts = urlsplit(_postgres_url(POSTGRES_URL))
    scheme = "postgresql+asyncpg" if parts.scheme in {"postgres", "postgresql"} else parts.scheme
    url = urlunsplit(parts._replace(scheme=scheme))
    schema = f"scheduler_owner_cap_{uuid.uuid4().hex}"
    engines = [create_async_engine(url, connect_args=build_asyncpg_connect_args(schema)) for _ in range(count)]
    try:
        async with engines[0].begin() as conn:
            await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
            await conn.run_sync(
                lambda sync: Base.metadata.create_all(
                    sync,
                    tables=[
                        ScheduledTaskRow.__table__,
                        ScheduledTaskRunRow.__table__,
                        RunRow.__table__,
                        RunChangeClockRow.__table__,
                        RunEventRow.__table__,
                        RunEventThreadSeqRow.__table__,
                    ],
                )
            )
        yield [async_sessionmaker(engine, expire_on_commit=False) for engine in engines]
    finally:
        async with engines[0].begin() as conn:
            await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
        for engine in engines:
            await engine.dispose()


@pytest.mark.asyncio
@pytest.mark.parametrize("global_cap", [3, 1], ids=["owner-cap-binds", "global-cap-binds"])
async def test_postgres_per_owner_cap_is_atomic_under_concurrent_claims(global_cap):
    per_user = 1
    async with _contender_engines(4) as factories:
        task_repo = ScheduledTaskRepository(factories[0])
        contenders = [ScheduledTaskRunRepository(sf) for sf in factories]
        seed = contenders[0]
        now = datetime.now(UTC)
        for round_index in range(20):
            rows = {f"a1-{round_index}": "owner-a", f"a2-{round_index}": "owner-a", f"b-{round_index}": "owner-b"}
            for row_id, owner in rows.items():
                await task_repo.create(
                    task_id=f"task-{row_id}",
                    user_id=owner,
                    thread_id=None,
                    context_mode="fresh_thread_per_run",
                    assistant_id="lead_agent",
                    title=row_id,
                    prompt="p",
                    schedule_type="cron",
                    schedule_spec={"cron": "* * * * *"},
                    timezone="UTC",
                    next_run_at=None,
                )
                await seed.create(run_record_id=row_id, task_id=f"task-{row_id}", thread_id=f"thread-{row_id}", scheduled_for=now, trigger="scheduled", status="queued")
            # Four contenders on separate engines: both A rows, B's row, and a
            # second claimer racing for A's first row.
            targets = [f"a1-{round_index}", f"a2-{round_index}", f"b-{round_index}", f"a1-{round_index}"]
            await asyncio.gather(
                *(
                    repo.claim_queued_run(target, lease_owner=f"pod-{index}", now=now, lease_seconds=60, global_max_concurrent_runs=global_cap, per_user_max_concurrent_runs=per_user)
                    for index, (repo, target) in enumerate(zip(contenders, targets, strict=True))
                )
            )
            async with factories[0]() as session:
                launching = {
                    row_id: owner
                    for row_id, owner in (
                        await session.execute(select(ScheduledTaskRunRow.id, ScheduledTaskRow.user_id).join(ScheduledTaskRow, ScheduledTaskRow.id == ScheduledTaskRunRow.task_id).where(ScheduledTaskRunRow.status == "launching"))
                    ).all()
                }
            owners = list(launching.values())
            assert all(owners.count(owner) <= per_user for owner in set(owners)), launching
            if global_cap >= 2:
                assert sorted(owners) == ["owner-a", "owner-b"], launching
            else:
                assert len(owners) == 1, launching
            # Free the slots before the next round.
            async with factories[0]() as session:
                await session.execute(update(ScheduledTaskRunRow).where(ScheduledTaskRunRow.status.in_(("queued", "launching"))).values(status="success", lease_owner=None, lease_expires_at=None))
                await session.commit()