agentci/tests/test_repository_tasks.py

import sqlite3
from contextlib import closing
from datetime import UTC, datetime, timedelta
from pathlib import Path

import pytest

from agentci.engine import _sqlite
from agentci.engine import repository as repository_module
from agentci.engine.events import JobStarted, PermissionGranted
from agentci.engine.model import IncomingCommand, QueueName, Task, TaskKind
from agentci.engine.repository import Repository

MIGRATIONS = Path(__file__).parents[1] / "src" / "agentci" / "migrations"


@pytest.fixture
async def repository(tmp_path: Path) -> Repository:
    value = Repository(tmp_path / "state.sqlite3", MIGRATIONS)
    await value.initialize()
    return value


def command(delivery: str, *, issue: int = 3) -> IncomingCommand:
    return IncomingCommand(
        delivery_id=delivery,
        comment_id=int(delivery.rsplit("-", 1)[-1]),
        repo_owner="alice",
        repo_name="repo",
        issue_number=issue,
        pr_number=None,
        requester="alice",
        body="/agent plan",
    )


def task_storage(repository: Repository, task_id: int) -> tuple[object, ...]:
    with closing(sqlite3.connect(repository.database_path)) as connection, connection:
        row = connection.execute(
            """SELECT status, attempts, available_at, error, created_at, started_at, finished_at
            FROM listener_tasks WHERE id=?""",
            (task_id,),
        ).fetchone()
    assert row is not None
    return row


async def test_claim_and_complete_record_attempt_and_lifecycle_timestamps(
    repository: Repository,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    clock = {"now": "2026-02-01T00:00:00+00:00"}
    monkeypatch.setattr(_sqlite, "now", lambda: clock["now"])
    job = (await repository.accept(command("delivery-1"))).job
    clock["now"] = "2026-02-01T00:01:00+00:00"

    task = await repository.claim_task(QueueName.CONTROL)
    assert task is not None
    clock["now"] = "2026-02-01T00:02:00+00:00"
    await repository.complete_task(task.id)

    assert (task, task_storage(repository, task.id)) == (
        Task(
            id=task.id,
            job_id=job.id,
            source_event_id="delivery:delivery-1",
            kind=TaskKind.AUTHORIZE,
            queue=QueueName.CONTROL,
            attempts=1,
        ),
        (
            "completed",
            1,
            "2026-02-01T00:00:00+00:00",
            None,
            "2026-02-01T00:00:00+00:00",
            "2026-02-01T00:01:00+00:00",
            "2026-02-01T00:02:00+00:00",
        ),
    )


@pytest.mark.parametrize(("attempts", "delay_seconds"), [(1, 2), (9, 256)])
async def test_retry_uses_bounded_backoff_and_preserves_attempt_history_until_due(
    repository: Repository,
    monkeypatch: pytest.MonkeyPatch,
    attempts: int,
    delay_seconds: int,
) -> None:
    clock = {"now": "2026-02-01T00:00:00+00:00"}
    monkeypatch.setattr(_sqlite, "now", lambda: clock["now"])
    await repository.accept(command("delivery-1"))
    task = await repository.claim_task(QueueName.CONTROL)
    assert task is not None
    with closing(sqlite3.connect(repository.database_path)) as connection, connection:
        connection.execute("UPDATE listener_tasks SET attempts=? WHERE id=?", (attempts, task.id))

    retry_time = datetime(2026, 2, 1, 1, tzinfo=UTC)
    available_at = retry_time + timedelta(seconds=delay_seconds)

    class FixedDateTime:
        @staticmethod
        def now(_timezone: object) -> datetime:
            return retry_time

    monkeypatch.setattr(repository_module, "datetime", FixedDateTime)
    error = "failure: " + "x" * 1100
    await repository.retry_task(task.id, attempts, error)
    pending_storage = task_storage(repository, task.id)
    clock["now"] = (available_at - timedelta(microseconds=1)).isoformat()
    early_claim = await repository.claim_task(QueueName.CONTROL)
    clock["now"] = available_at.isoformat()
    due_claim = await repository.claim_task(QueueName.CONTROL)

    assert (
        pending_storage,
        early_claim,
        due_claim and due_claim.attempts,
    ) == (
        (
            "pending",
            attempts,
            available_at.isoformat(),
            error[:1000],
            "2026-02-01T00:00:00+00:00",
            "2026-02-01T00:00:00+00:00",
            None,
        ),
        None,
        attempts + 1,
    )


async def test_recovery_requeues_control_and_unstarted_execution_but_fails_started_execution(
    repository: Repository,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    clock = {"now": "2026-02-01T00:00:00+00:00"}
    monkeypatch.setattr(_sqlite, "now", lambda: clock["now"])

    await repository.accept(command("delivery-1", issue=1))
    control_task = await repository.claim_task(QueueName.CONTROL)
    assert control_task is not None

    queued_job = (await repository.accept(command("delivery-2", issue=2))).job
    await repository.apply("queued:grant", PermissionGranted(job_id=queued_job.id))
    queued_execute = await repository.claim_task(QueueName.JOBS)
    assert queued_execute is not None

    running_job = (await repository.accept(command("delivery-3", issue=3))).job
    await repository.apply("running:grant", PermissionGranted(job_id=running_job.id))
    running_execute = await repository.claim_task(QueueName.JOBS)
    assert running_execute is not None
    await repository.apply("running:start", JobStarted(job_id=running_job.id))

    clock["now"] = "2026-02-01T01:00:00+00:00"
    await repository.recover_tasks()

    assert (
        task_storage(repository, control_task.id),
        task_storage(repository, queued_execute.id),
        task_storage(repository, running_execute.id),
        [job.id for job in await repository.running_jobs()],
    ) == (
        (
            "pending",
            1,
            "2026-02-01T00:00:00+00:00",
            None,
            "2026-02-01T00:00:00+00:00",
            None,
            None,
        ),
        (
            "pending",
            1,
            "2026-02-01T00:00:00+00:00",
            None,
            "2026-02-01T00:00:00+00:00",
            None,
            None,
        ),
        (
            "failed",
            1,
            "2026-02-01T00:00:00+00:00",
            "Service restarted after execution began",
            "2026-02-01T00:00:00+00:00",
            "2026-02-01T00:00:00+00:00",
            "2026-02-01T01:00:00+00:00",
        ),
        [running_job.id],
    )