deer-flow/backend/tests/test_scheduler_occurrence_ordering.py

"""Late launches must preserve the latest occurrence's parent projection.

The scheduler and repositories are real. Only the external run lifecycle and
barriers between its launch return and scheduler bookkeeping are controlled.
"""

from __future__ import annotations

import asyncio
import os
import uuid
from contextlib import suppress
from datetime import UTC, datetime, timedelta
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit

import pytest
import pytest_asyncio
from sqlalchemy import event, select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.schema import CreateSchema, DropSchema

from app.scheduler.service import ScheduledTaskService
from deerflow.persistence.base import Base
from deerflow.persistence.run import RunRepository
from deerflow.persistence.run.model import 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.runs.manager import RunRecord
from deerflow.runtime.runs.schemas import DisconnectMode, RunStatus

pytestmark = pytest.mark.asyncio


def _asyncpg_url(url: str | None) -> str | None:
    """CI passes a sync ``postgresql://...?sslmode=disable`` URL; use asyncpg and drop libpq-only keys."""
    if not url:
        return url
    parts = urlsplit(url)
    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"}])
    return urlunsplit(parts._replace(scheme=scheme, query=query))


@pytest_asyncio.fixture(params=["sqlite", "postgres-single", "postgres-multi"])
async def occurrence_databases(request, tmp_path):
    backend = request.param
    postgres_uri = _asyncpg_url(os.environ.get("TEST_POSTGRES_URI"))
    if backend != "sqlite" and not postgres_uri:
        pytest.skip("TEST_POSTGRES_URI is not set")
    schema = "scheduler_order_" + uuid.uuid4().hex
    admin = None
    engines = []
    try:
        if backend != "sqlite":
            admin = create_async_engine(postgres_uri)
            async with admin.begin() as connection:
                await connection.execute(CreateSchema(schema))
        for _ in range(3):
            if backend == "sqlite":
                engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'scheduler.db'}")

                @event.listens_for(engine.sync_engine, "connect")
                def configure_sqlite(connection, _record):
                    cursor = connection.cursor()
                    cursor.execute("PRAGMA journal_mode=WAL")
                    cursor.execute("PRAGMA busy_timeout=10000")
                    cursor.execute("PRAGMA foreign_keys=ON")
                    cursor.close()
            else:
                engine = create_async_engine(postgres_uri, connect_args={"server_settings": {"search_path": schema}})
            engines.append(engine)
        async with engines[0].begin() as connection:
            await connection.run_sync(lambda sync: Base.metadata.create_all(sync, tables=[ScheduledTaskRow.__table__, ScheduledTaskRunRow.__table__, RunRow.__table__]))
        yield [async_sessionmaker(engine, expire_on_commit=False) for engine in engines], backend == "postgres-multi"
    finally:
        for engine in engines:
            await engine.dispose()
        if admin is not None:
            async with admin.begin() as connection:
                await connection.execute(DropSchema(schema, cascade=True, if_exists=True))
            await admin.dispose()


class _RevivalBarrierRepository(ScheduledTaskRunRepository):
    def __init__(self, factory, durable, revived, resume):
        super().__init__(factory, run_repository=durable)
        self.revived = revived
        self.resume = resume

    async def reconcile_launched_run(self, *args, **kwargs):
        result = await super().reconcile_launched_run(*args, **kwargs)
        assert result is True
        self.revived.set()
        await self.resume.wait()
        return result


def _service(tasks, occurrences, launcher, multi):
    return ScheduledTaskService(
        task_repo=tasks,
        task_run_repo=occurrences,
        launch_run=launcher,
        poll_interval_seconds=3600,
        lease_seconds=5,
        max_concurrent_runs=3,
        queue_timeout_seconds=10,
        multi_instance=multi,
        run_lease_grace_seconds=0,
    )


async def _parent_projection(factory, task_id):
    async with factory() as session:
        parent = await session.get(ScheduledTaskRow, task_id)
        return {key: getattr(parent, key) for key in ("status", "last_error", "last_run_id", "last_thread_id", "last_run_at", "next_run_at", "lease_owner", "lease_expires_at", "run_count")}


def _completion(run_id, launch_args, status, error=None):
    return RunRecord(
        run_id=run_id,
        thread_id=launch_args["thread_id"],
        assistant_id=None,
        status=status,
        on_disconnect=DisconnectMode.continue_,
        user_id="user-order",
        metadata=launch_args["metadata"],
        error=error,
    )


@pytest.mark.parametrize("interrupt_before_parent", [True, False], ids=["restart-after-revival", "late-completion-callback"])
@pytest.mark.parametrize("newer_occurrence", [True, False], ids=["newer-failed-occurrence", "single-occurrence-control"])
async def test_late_launch_preserves_newer_result_and_counts_once(occurrence_databases, interrupt_before_parent, newer_occurrence):
    factories, multi = occurrence_databases
    durable_a, durable_b, durable_recovery = [RunRepository(factory) for factory in factories]
    tasks_a, tasks_b, tasks_recovery = [ScheduledTaskRepository(factory, run_repository=durable) for factory, durable in zip(factories, (durable_a, durable_b, durable_recovery), strict=True)]
    entered, release, revived, resume = [asyncio.Event() for _ in range(4)]
    occurrences_a = _RevivalBarrierRepository(factories[0], durable_a, revived, resume)
    occurrences_b = ScheduledTaskRunRepository(factories[1], run_repository=durable_b)
    occurrences_recovery = ScheduledTaskRunRepository(factories[2], run_repository=durable_recovery)
    now = datetime.now(UTC)
    task_id, run_a, run_b = "task-order", "run-order-a", "run-order-b"
    launches = {}

    async def launch_a(**kwargs):
        launches["a"] = kwargs
        entered.set()
        await release.wait()
        # Model a real run becoming terminal before its slow launch call
        # returns. Scheduled occurrence states are only changed by the service.
        await durable_a.put(run_a, thread_id=kwargs["thread_id"], user_id="user-order", status="running", metadata=kwargs["metadata"])
        await durable_a.update_status(run_a, "success")
        return {"run_id": run_a, "thread_id": kwargs["thread_id"]}

    async def launch_b(**kwargs):
        launches["b"] = kwargs
        await durable_b.put(run_b, thread_id=kwargs["thread_id"], user_id="user-order", status="running", metadata=kwargs["metadata"])
        return {"run_id": run_b, "thread_id": kwargs["thread_id"]}

    async def singleton_launcher(**kwargs):
        if "a" not in launches:
            return await launch_a(**kwargs)
        assert "b" not in launches
        return await launch_b(**kwargs)

    async def unexpected_launch(**_kwargs):
        pytest.fail("recovery must not launch another occurrence")

    service_a = _service(tasks_a, occurrences_a, launch_a if multi else singleton_launcher, multi)
    service_b = _service(tasks_b, occurrences_b, launch_b, multi) if multi else service_a
    peer = _service(tasks_recovery, occurrences_recovery, unexpected_launch, multi) if multi else service_a
    scheduled_at = now - timedelta(seconds=1) if multi else now + timedelta(days=1)
    await tasks_a.create(
        task_id=task_id,
        user_id="user-order",
        thread_id=None,
        context_mode="fresh_thread_per_run",
        assistant_id=None,
        title="occurrence order",
        prompt="test",
        schedule_type="once",
        schedule_spec={"run_at": scheduled_at.isoformat()},
        timezone="UTC",
        next_run_at=scheduled_at,
    )
    initial_task = await tasks_a.get(task_id, user_id="user-order")
    dispatch = asyncio.create_task(service_a.run_once(now=now) if multi else service_a.dispatch_task(initial_task, now=now, trigger="manual"))
    restarted = _service(tasks_recovery, occurrences_recovery, unexpected_launch, multi)

    async def parked_poll_loop():
        await restarted._stop.wait()

    restarted._run_loop = parked_poll_loop
    try:
        await asyncio.wait_for(entered.wait(), 10)
        assert (await occurrences_b.list_by_task(task_id))[0]["status"] == "launching"
        # Explicit scheduler time expires the launch claim and then the queue;
        # no sleeps or direct terminal occurrence writes are needed.
        await peer.run_once(now=now + timedelta(seconds=30))
        expired = (await occurrences_b.list_by_task(task_id))[0]
        assert expired["status"] == "failed"
        assert expired["run_id"] is None

        if newer_occurrence:
            result = await service_b.dispatch_task(await tasks_b.get(task_id, user_id="user-order"), now=now + timedelta(seconds=31), trigger="manual")
            assert result["outcome"] == "launched"
            await durable_b.update_status(run_b, "error", error="newer run failed")
            await service_b.handle_run_completion(_completion(run_b, launches["b"], RunStatus.error, "newer run failed"))
            expected = await _parent_projection(factories[1], task_id)
            assert expected["status"] == "failed"
            assert expected["last_run_id"] == run_b
            assert expected["run_count"] == 1
            expected["run_count"] = 2

        release.set()
        await asyncio.wait_for(revived.wait(), 10)
        if interrupt_before_parent:
            # The revival transaction committed, but dispatch has not yet
            # performed the parent launch accounting.
            dispatch.cancel()
            with suppress(asyncio.CancelledError):
                await dispatch
        else:
            resume.set()
            await asyncio.wait_for(dispatch, 10)
            await service_a.handle_run_completion(_completion(run_a, launches["a"], RunStatus.success))
            if newer_occurrence:
                assert await _parent_projection(factories[1], task_id) == expected

        await restarted.start()
        await restarted.stop()
        after = await _parent_projection(factories[1], task_id)
        if newer_occurrence:
            assert after == expected
        else:
            assert after["status"] == "completed"
            assert after["last_run_id"] == run_a
            assert after["run_count"] == 1
        async with factories[1]() as session:
            statuses = dict((await session.execute(select(ScheduledTaskRunRow.run_id, ScheduledTaskRunRow.status).where(ScheduledTaskRunRow.task_id == task_id))).all())
        assert statuses == ({run_a: "success", run_b: "failed"} if newer_occurrence else {run_a: "success"})

        # Repeated completion and recovery must neither re-project old state
        # nor account either successful launch a second time.
        await service_a.handle_run_completion(_completion(run_a, launches["a"], RunStatus.success))
        assert await _parent_projection(factories[1], task_id) == after
        if newer_occurrence:
            await service_b.handle_run_completion(_completion(run_b, launches["b"], RunStatus.error, "newer run failed"))
            assert await _parent_projection(factories[1], task_id) == after
        await restarted.start()
        await restarted.stop()
        assert await _parent_projection(factories[1], task_id) == after
    finally:
        release.set()
        resume.set()
        if not dispatch.done():
            dispatch.cancel()
        with suppress(asyncio.CancelledError):
            await dispatch
        await restarted.stop()


@pytest.mark.parametrize("interrupt_before_return", [True, False], ids=["crash-before-launch-return", "normal-launch-return"])
async def test_completion_before_launch_return_accounts_once(occurrence_databases, interrupt_before_return):
    factories, multi = occurrence_databases
    durable = RunRepository(factories[0])
    tasks = ScheduledTaskRepository(factories[0], run_repository=durable)
    occurrences = ScheduledTaskRunRepository(factories[0], run_repository=durable)
    completed, release = asyncio.Event(), asyncio.Event()
    run_id, task_id = "run-fast-completion", "task-fast-completion"
    now = datetime.now(UTC)
    record = None

    async def launcher(**kwargs):
        nonlocal record
        await durable.put(run_id, thread_id=kwargs["thread_id"], user_id="user-order", status="running", metadata=kwargs["metadata"])
        await durable.update_status(run_id, "success")
        record = _completion(run_id, kwargs, RunStatus.success)
        await service.handle_run_completion(record)
        completed.set()
        await release.wait()
        return {"run_id": run_id, "thread_id": kwargs["thread_id"]}

    async def unexpected_launch(**_kwargs):
        pytest.fail("completed occurrence must not launch again")

    service = _service(tasks, occurrences, launcher, multi)
    await tasks.create(
        task_id=task_id,
        user_id="user-order",
        thread_id=None,
        context_mode="fresh_thread_per_run",
        assistant_id=None,
        title="fast completion",
        prompt="test",
        schedule_type="once",
        schedule_spec={"run_at": now.isoformat()},
        timezone="UTC",
        next_run_at=now,
    )
    dispatch = asyncio.create_task(service.dispatch_task(await tasks.get(task_id, user_id="user-order"), now=now, trigger="manual"))
    recovery_durable = RunRepository(factories[1])
    restarted = _service(
        ScheduledTaskRepository(factories[1], run_repository=recovery_durable),
        ScheduledTaskRunRepository(factories[1], run_repository=recovery_durable),
        unexpected_launch,
        multi,
    )

    async def parked_poll_loop():
        await restarted._stop.wait()

    restarted._run_loop = parked_poll_loop
    try:
        await asyncio.wait_for(completed.wait(), 10)
        before_return = await _parent_projection(factories[1], task_id)
        assert before_return["status"] == "completed"
        assert before_return["run_count"] == 1
        assert before_return["last_run_id"] == run_id
        assert before_return["last_thread_id"] == record.thread_id
        if interrupt_before_return:
            dispatch.cancel()
            with suppress(asyncio.CancelledError):
                await dispatch
        else:
            release.set()
            await asyncio.wait_for(dispatch, 10)

        await restarted.start()
        await restarted.stop()
        after = await _parent_projection(factories[1], task_id)
        assert after["status"] == "completed"
        assert after["last_run_id"] == run_id
        assert after["run_count"] == 1
        assert (await occurrences.list_by_task(task_id))[0]["status"] == "success"
        await service.handle_run_completion(record)
        await restarted.start()
        await restarted.stop()
        assert await _parent_projection(factories[1], task_id) == after
    finally:
        release.set()
        if not dispatch.done():
            dispatch.cancel()
        with suppress(asyncio.CancelledError):
            await dispatch
        await restarted.stop()