deer-flow/backend/tests/test_scheduled_occurrence_sequence.py

"""Database ordering is per task and survives retries independently of caller clocks."""

from __future__ import annotations

import asyncio
import os
import uuid
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
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

import deerflow.persistence.models  # noqa: F401
from deerflow.persistence.base import Base
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
from deerflow.persistence.scheduled_task_runs import ActiveScheduledRunConflict, ScheduledTaskRunRepository
from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
from deerflow.persistence.scheduled_tasks.model import ACTIVE_RUN_STATUSES, ONCE_TASK_STATUS_BY_RUN_STATUS, ScheduledTaskRow

pytestmark = pytest.mark.asyncio


@pytest_asyncio.fixture(params=["sqlite", "postgres"])
async def occurrence_factories(request, tmp_path):
    """Two pools guarantee competing admissions use independent DB connections."""
    schema = None
    if request.param == "postgres":
        uri = os.environ.get("TEST_POSTGRES_URI")
        if not uri:
            pytest.skip("requires TEST_POSTGRES_URI (real Postgres for occurrence ordering)")
        parts = urlsplit(uri)
        # CI passes a sync ``postgresql://...?sslmode=disable`` URL; the async
        # engine needs the asyncpg driver and rejects libpq-only query keys.
        scheme = "postgresql+asyncpg" if parts.scheme in {"postgres", "postgresql"} else parts.scheme
        query = urlencode([(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key not in {"sslmode", "channel_binding"}])
        uri = urlunsplit(parts._replace(scheme=scheme, query=query))
        schema = f"occurrence_{uuid.uuid4().hex}"
        options = {"connect_args": build_asyncpg_connect_args(schema)}
    else:
        uri = f"sqlite+aiosqlite:///{tmp_path / 'occurrences.db'}"
        options = {"connect_args": {"timeout": 30}}
    engines = [create_async_engine(uri, **options) for _ in range(2)]
    try:
        async with engines[0].begin() as connection:
            if schema:
                await connection.execute(text(f'CREATE SCHEMA "{schema}"'))
            await connection.run_sync(Base.metadata.create_all)
        yield tuple(async_sessionmaker(engine, expire_on_commit=False) for engine in engines)
    finally:
        if schema:
            async with engines[0].begin() as connection:
                await connection.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
        for engine in engines:
            await engine.dispose()


async def _create_task(factory, task_id="task", *, schedule_type="cron"):
    spec = {"cron": "* * * * *"} if schedule_type == "cron" else {"run_at": datetime(2026, 7, 15, 12, 0, tzinfo=UTC).isoformat()}
    return await ScheduledTaskRepository(factory).create(
        task_id=task_id,
        user_id="user-1",
        thread_id="thread-1",
        context_mode="reuse_thread",
        assistant_id=None,
        title="Occurrence ordering",
        prompt="p",
        schedule_type=schedule_type,
        schedule_spec=spec,
        timezone="UTC",
        next_run_at=None,
    )


async def _create_run(factory, run_id, *, task_id="task", status="success"):
    return await ScheduledTaskRunRepository(factory).create(
        run_record_id=run_id,
        task_id=task_id,
        thread_id=f"thread-{run_id}",
        scheduled_for=datetime(2026, 7, 15, 12, 0, tzinfo=UTC),
        trigger="manual",
        status=status,
    )


async def _sequence(factory, run_id):
    async with factory() as session:
        return await session.scalar(select(ScheduledTaskRunRow.occurrence_seq).where(ScheduledTaskRunRow.id == run_id))


async def _high_water_mark(factory, task_id="task"):
    async with factory() as session:
        return await session.scalar(select(ScheduledTaskRow.last_occurrence_seq).where(ScheduledTaskRow.id == task_id))


async def test_concurrent_allocations_use_distinct_monotonic_sequences(occurrence_factories):
    first, second = occurrence_factories
    original = await _create_task(first)
    ready = [asyncio.Event(), asyncio.Event()]
    start = asyncio.Event()

    async def admit(factory, lane):
        ready[lane].set()
        await start.wait()
        for index in range(4):
            await _create_run(factory, f"run-{lane}-{index}")

    admissions = [asyncio.create_task(admit(factory, lane)) for lane, factory in enumerate((first, second))]
    await asyncio.gather(*(event.wait() for event in ready))
    start.set()
    await asyncio.gather(*admissions)
    sequences = [await _sequence(first, f"run-{lane}-{index}") for lane in range(2) for index in range(4)]
    assert sorted(sequences) == list(range(1, 9))
    for lane in range(2):
        lane_sequences = sequences[lane * 4 : (lane + 1) * 4]
        assert lane_sequences == sorted(lane_sequences)
    assert await _high_water_mark(first) == 8
    current = await ScheduledTaskRepository(first).get("task", user_id="user-1")
    assert current["updated_at"] == original["updated_at"]
    assert current["run_count"] == 0


async def test_sequence_allocation_is_independent_per_task(occurrence_factories):
    first, second = occurrence_factories
    for task_id in ("task-a", "task-b"):
        await _create_task(first, task_id)
    await _create_run(first, "run-a1", task_id="task-a")
    await _create_run(second, "run-a2", task_id="task-a")
    await _create_run(second, "run-b1", task_id="task-b")
    assert [await _sequence(first, run_id) for run_id in ("run-a1", "run-a2", "run-b1")] == [1, 2, 1]


async def test_active_conflict_rolls_back_sequence_allocation(occurrence_factories):
    first, second = occurrence_factories
    await _create_task(first)
    await _create_run(first, "active", status="queued")
    with pytest.raises(ActiveScheduledRunConflict):
        await _create_run(second, "rejected", status="queued")
    assert await _high_water_mark(first) == 1
    assert await _sequence(first, "rejected") is None
    await ScheduledTaskRunRepository(first).update_status("active", status="success")
    await _create_run(second, "accepted", status="queued")
    assert await _sequence(first, "accepted") == 2


@pytest.mark.parametrize("status", ["queued", "success"])
async def test_primary_key_conflict_is_not_an_active_conflict_and_rolls_back(occurrence_factories, status):
    first, second = occurrence_factories
    await _create_task(first, "task-a")
    await _create_task(first, "task-b")
    await _create_run(first, "duplicate", task_id="task-a")
    with pytest.raises(IntegrityError):
        await _create_run(second, "duplicate", task_id="task-b", status=status)
    assert await _high_water_mark(first, "task-b") == 0
    await _create_run(second, "unique", task_id="task-b", status=status)
    assert await _sequence(first, "unique") == 1


async def test_requeue_and_reclaim_preserve_occurrence_sequence(occurrence_factories):
    first, second = occurrence_factories
    await _create_task(first)
    await _create_run(first, "retry", status="queued")
    now = datetime(2026, 7, 15, 12, 0, tzinfo=UTC)
    repo = ScheduledTaskRunRepository(second)
    for attempt in range(2):
        claimed = await repo.claim_queued_run("retry", now=now, lease_owner="worker", lease_seconds=60, global_max_concurrent_runs=1)
        assert claimed is not None
        assert claimed["attempt_count"] == attempt + 1
        assert await repo.requeue_claimed_run("retry", lease_owner="worker") is True
    assert await _sequence(first, "retry") == 1
    assert await _high_water_mark(first) == 1


async def test_internal_sequence_fields_are_absent_from_repository_responses(occurrence_factories):
    first, _second = occurrence_factories
    created_task = await _create_task(first)
    task_repo = ScheduledTaskRepository(first)
    run_repo = ScheduledTaskRunRepository(first)
    created_run = await _create_run(first, "queued", status="queued")
    task_responses = [created_task, await task_repo.get("task", user_id="user-1"), *(await task_repo.list_by_user("user-1"))]
    run_responses = [created_run, await run_repo.get_active_run("task"), *(await run_repo.list_by_task("task")), *(await run_repo.list_queued_runs(limit=10))]
    for response in task_responses + run_responses:
        assert {"last_occurrence_seq", "occurrence_seq", "launch_accounted"}.isdisjoint(response)
    assert await _sequence(first, "queued") == 1


async def _insert_unsequenced_run(factory, run_id, *, created_at, status="success"):
    """Insert without the repository: legacy history or a pre-upgrade writer."""
    async with factory() as session:
        session.add(
            ScheduledTaskRunRow(
                id=run_id,
                task_id="task",
                thread_id=f"thread-{run_id}",
                scheduled_for=created_at,
                created_at=created_at,
                trigger="manual",
                status=status,
            )
        )
        await session.commit()


async def _latest_run_id(factory):
    async with factory() as session:
        latest = await ScheduledTaskRepository._fetch_latest_run(session, "task")
        assert latest is not None
        return latest.id


async def test_recovery_order_keeps_timestamp_fallback_for_legacy_only_history(occurrence_factories):
    first, _second = occurrence_factories
    await _create_task(first)
    now = datetime(2026, 7, 15, 12, 0, tzinfo=UTC)
    # Unsequenced history is deliberately not assigned guessed sequence values.
    for index in range(2):
        await _insert_unsequenced_run(first, f"legacy-{index}", created_at=now + timedelta(days=index))
    assert await _latest_run_id(first) == "legacy-1"
    assert await _sequence(first, "legacy-0") is None
    assert await _sequence(first, "legacy-1") is None
    assert await _high_water_mark(first) == 0


@pytest.mark.parametrize("unsequenced_status", ["skipped", "running"])
@pytest.mark.parametrize("unsequenced_offset", [timedelta(days=-365), timedelta(seconds=30)], ids=["legacy-history-is-older", "pre-upgrade-writer-is-newer"])
async def test_recovery_lookup_prefers_the_highest_sequence_whenever_one_exists(occurrence_factories, unsequenced_offset, unsequenced_status):
    """Sequence decides whenever a sequenced row exists.

    Reversed caller clocks between sequenced rows do not matter, and an
    unsequenced row (legacy history or a pre-upgrade Gateway writer) is not
    consulted even when its caller timestamp is later: the lookup returns the
    same row ``can_project`` accepts, so recovery cannot act on a row that the
    other parent writes would reject.
    """
    first, second = occurrence_factories
    await _create_task(first)
    now = datetime(2026, 7, 15, 12, 0, tzinfo=UTC)
    # The later sequence carries the earlier caller clock: sequence still wins.
    with patch("deerflow.persistence.scheduled_task_runs.sql.datetime") as clock:
        clock.now.return_value = now + timedelta(seconds=30)
        await _create_run(first, "older-sequenced")
        clock.now.return_value = now
        await _create_run(second, "newer-sequenced")
    await _insert_unsequenced_run(first, "unsequenced", created_at=now + unsequenced_offset, status=unsequenced_status)
    assert [await _sequence(first, run_id) for run_id in ("older-sequenced", "newer-sequenced", "unsequenced")] == [1, 2, None]
    assert await _latest_run_id(first) == "newer-sequenced"


@pytest.mark.parametrize("recovery_method", ["cancel_stuck_once_tasks", "reconcile_stuck_once_tasks"])
@pytest.mark.parametrize(
    ("sequenced_status", "unsequenced_status"),
    [("running", "skipped"), ("success", "running"), ("success", "interrupted"), ("failed", "running")],
    ids=[
        "unsequenced-skipped-while-sequenced-active",
        "unsequenced-running-while-sequenced-success",
        "unsequenced-interrupted-while-sequenced-success",
        "unsequenced-running-while-sequenced-failed",
    ],
)
async def test_once_recovery_defers_while_any_occurrence_is_live_then_projects_the_sequence_winner(occurrence_factories, recovery_method, sequenced_status, unsequenced_status):
    """Mixed-writer interleavings from review, both recovery paths, both backends.

    A live occurrence row is the task's newest admission by construction
    (``uq_scheduled_task_run_active``), whatever its caller clock and whether or
    not it carries a sequence, so recovery defers while one exists: an
    unsequenced ``skipped`` row cannot cancel a parent whose sequenced
    occurrence is live, and an unsequenced ``running`` row cannot be skipped
    over to finalise the parent from an older sequenced outcome.  Once no row is
    live, the sequence winner decides and a terminalised unsequenced row never
    overrides it.
    """
    first, _second = occurrence_factories
    await _create_task(first, schedule_type="once")
    task_repo = ScheduledTaskRepository(first)
    await task_repo.update("task", user_id="user-1", updates={"status": "running"})
    await _create_run(first, "sequenced", status=sequenced_status)
    # A pre-upgrade node on a skewed clock: no sequence, later caller timestamp.
    await _insert_unsequenced_run(first, "unsequenced", created_at=datetime.now(UTC) + timedelta(minutes=5), status=unsequenced_status)
    assert await _sequence(first, "sequenced") == 1
    assert await _high_water_mark(first) == 1

    kwargs = {"error": "interrupted: recovery"}
    if recovery_method == "reconcile_stuck_once_tasks":
        kwargs["now"] = datetime.now(UTC) + timedelta(minutes=10)

    async def recover():
        count = await getattr(task_repo, recovery_method)(**kwargs)
        task = await task_repo.get_internal("task")
        assert task is not None
        return count, task

    any_live = sequenced_status in ACTIVE_RUN_STATUSES or unsequenced_status in ACTIVE_RUN_STATUSES
    sequence_outcome = ONCE_TASK_STATUS_BY_RUN_STATUS.get(sequenced_status, "running")
    expected_status, expected_count = ("running", 0) if any_live else (sequence_outcome, 1)
    for _ in range(2):  # a second pass must not change the outcome
        count, task = await recover()
        assert task["status"] == expected_status
        assert task["last_error"] is None
        assert count == expected_count
        expected_count = 0

    if unsequenced_status in ACTIVE_RUN_STATUSES:
        # The pre-upgrade node died and occurrence recovery terminalised its
        # row: the sequence winner now decides, not the newer unsequenced row.
        assert await ScheduledTaskRunRepository(first).update_status("unsequenced", status="interrupted", error="pre-upgrade node died")
        count, task = await recover()
        assert task["status"] == sequence_outcome
        assert task["last_error"] is None
        assert count == 1