agentci/tests/test_worker.py

import asyncio
from dataclasses import replace
from pathlib import Path
from types import SimpleNamespace

import pytest

from agentci.engine.events import (
    CommentLinked,
    JobCompleted,
    JobFailed,
    JobStarted,
    PermissionDenied,
    PermissionGranted,
    ServiceRestarted,
)
from agentci.engine.events import (
    JobRejected as RejectedEvent,
)
from agentci.engine.model import (
    IncomingCommand,
    Job,
    JobKind,
    JobStatus,
    QueueName,
    Task,
    TaskKind,
    Workflow,
    WorkflowKind,
)
from agentci.engine.reducer import render_job_comment
from agentci.engine.repository import Repository
from agentci.gitea import CommentInfo
from agentci.worker import Worker, _safe_error
from agentci.workflows.render import JobRejected

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


class FakeRepository:
    def __init__(
        self,
        value: Job | None = None,
        *,
        workflow: Workflow | None = None,
    ) -> None:
        self.job = value
        self.workflow = workflow
        self.tasks: list[Task] = []
        self.claimed_queues: list[QueueName] = []
        self.completed: list[int] = []
        self.retried: list[tuple[int, int, str]] = []
        self.events: list[tuple[str, object]] = []
        self.operations: list[str] = []
        self.running: list[Job] = []
        self.failed_workflows: list[str] = []
        self.complete_stop: asyncio.Event | None = None

    async def claim_task(self, queue: QueueName) -> Task | None:
        self.claimed_queues.append(queue)
        return self.tasks.pop(0) if self.tasks else None

    async def complete_task(self, task_id: int) -> None:
        self.completed.append(task_id)
        if self.complete_stop is not None:
            self.complete_stop.set()

    async def retry_task(self, task_id: int, attempts: int, error: str) -> None:
        self.retried.append((task_id, attempts, error))

    async def get_job(self, _job_id: str) -> Job | None:
        return self.job

    async def apply(self, event_id: str, event: object) -> SimpleNamespace:
        self.operations.append(f"apply:{event_id}")
        self.events.append((event_id, event))
        if isinstance(event, JobStarted) and self.job is not None:
            self.job = replace(self.job, status=JobStatus.RUNNING, stage="starting")
        elif isinstance(event, CommentLinked) and self.job is not None:
            self.job = replace(self.job, accepted_comment_id=event.comment_id)
        return SimpleNamespace(job=self.job)

    async def recover_tasks(self) -> None:
        self.operations.append("recover_tasks")

    async def running_jobs(self) -> list[Job]:
        self.operations.append("running_jobs")
        return self.running

    async def get_workflow(self, _workflow_id: str) -> Workflow | None:
        return self.workflow

    async def fail_job_workflow(self, job_id: str) -> None:
        self.failed_workflows.append(job_id)


class FakeGitea:
    def __init__(self, *, permitted: bool = True) -> None:
        self.permitted = permitted
        self.permission_calls: list[tuple[str, str, str]] = []
        self.update_results: list[bool] = [True]
        self.update_calls: list[tuple[str, str, int, str]] = []
        self.comments: list[CommentInfo] = []
        self.issue_comment_calls: list[tuple[str, str, int]] = []
        self.created_comment_id = 42
        self.create_calls: list[tuple[str, str, int, str]] = []

    async def has_write_permission(self, owner: str, repo: str, user: str) -> bool:
        self.permission_calls.append((owner, repo, user))
        return self.permitted

    async def update_comment(self, owner: str, repo: str, comment_id: int, body: str) -> bool:
        self.update_calls.append((owner, repo, comment_id, body))
        if len(self.update_results) > 1:
            return self.update_results.pop(0)
        return self.update_results[0]

    async def issue_comments(self, owner: str, repo: str, issue: int) -> list[CommentInfo]:
        self.issue_comment_calls.append((owner, repo, issue))
        return self.comments

    async def create_comment(self, owner: str, repo: str, issue: int, body: str) -> int:
        self.create_calls.append((owner, repo, issue, body))
        return self.created_comment_id


class FakeOpenCode:
    def __init__(self, ready: list[bool] | None = None) -> None:
        self.ready_results = ready or [True]
        self.ready_calls = 0
        self.aborted: list[tuple[str, Path]] = []

    async def ready(self) -> bool:
        self.ready_calls += 1
        if len(self.ready_results) > 1:
            return self.ready_results.pop(0)
        return self.ready_results[0]

    async def abort(self, session_id: str, workspace: Path) -> None:
        self.aborted.append((session_id, workspace))


def job(
    *,
    job_id: str = "job",
    status: JobStatus = JobStatus.RECEIVED,
    workflow_id: str | None = None,
    session_id: str | None = None,
    accepted_comment_id: int | None = None,
) -> Job:
    return Job(
        id=job_id,
        target_key="org/repo:issue:1",
        repo_owner="org",
        repo_name="repo",
        issue_number=1,
        pr_number=None,
        requester="alice",
        comment_id=1,
        delivery_id=f"delivery-{job_id}",
        receive_sequence=1,
        command_body="/agent plan",
        kind=JobKind.PLAN if status is not JobStatus.RECEIVED else None,
        status=status,
        stage=status.value,
        workflow_id=workflow_id,
        runtime_session_id=session_id,
        accepted_comment_id=accepted_comment_id,
    )


def task(
    kind: TaskKind = TaskKind.EXECUTE,
    queue: QueueName = QueueName.JOBS,
    *,
    task_id: int = 7,
) -> Task:
    return Task(task_id, "job", "source", kind, queue, 2)


def make_worker(
    tmp_path: Path,
    repository: FakeRepository | Repository,
    *,
    gitea: FakeGitea | None = None,
    opencode: FakeOpenCode | None = None,
    max_concurrent_jobs: int = 2,
) -> Worker:
    return Worker(
        repository=repository,  # type: ignore[arg-type]
        gitea=gitea or FakeGitea(),  # type: ignore[arg-type]
        opencode=opencode or FakeOpenCode(),  # type: ignore[arg-type]
        services=SimpleNamespace(),  # type: ignore[arg-type]
        poll_seconds=1,
        max_concurrent_jobs=max_concurrent_jobs,
        workspaces_dir=tmp_path,
        bot_username="agentci",
    )


async def test_loop_completes_successful_task(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    repository = FakeRepository()
    queued_task = task(queue=QueueName.CONTROL)
    repository.tasks = [queued_task]
    stop = asyncio.Event()
    handled: list[Task] = []
    value = make_worker(tmp_path, repository)

    async def handle(current: Task) -> None:
        handled.append(current)
        stop.set()

    monkeypatch.setattr(value, "_handle", handle)

    await value._loop(QueueName.CONTROL, stop)

    assert handled == [queued_task]
    assert repository.completed == [queued_task.id]
    assert repository.retried == []


async def test_loop_retries_failed_task(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    repository = FakeRepository()
    queued_task = task(queue=QueueName.CONTROL)
    repository.tasks = [queued_task]
    stop = asyncio.Event()
    value = make_worker(tmp_path, repository)

    async def fail(_current: Task) -> None:
        stop.set()
        raise RuntimeError("provider\nfailed")

    monkeypatch.setattr(value, "_handle", fail)

    await value._loop(QueueName.CONTROL, stop)

    assert repository.completed == []
    assert repository.retried == [
        (queued_task.id, queued_task.attempts, "RuntimeError: provider failed")
    ]


async def test_loop_propagates_cancellation_without_rescheduling(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    repository = FakeRepository()
    queued_task = task(queue=QueueName.CONTROL)
    repository.tasks = [queued_task]
    stop = asyncio.Event()
    value = make_worker(tmp_path, repository)

    async def cancel(_current: Task) -> None:
        stop.set()
        raise asyncio.CancelledError

    monkeypatch.setattr(value, "_handle", cancel)

    with pytest.raises(asyncio.CancelledError):
        await value._loop(QueueName.CONTROL, stop)

    assert repository.completed == []
    assert repository.retried == []


async def test_loop_completes_task_whose_job_no_longer_exists(tmp_path: Path) -> None:
    repository = FakeRepository(None)
    missing_job_task = task(queue=QueueName.CONTROL)
    repository.tasks = [missing_job_task]
    stop = asyncio.Event()
    repository.complete_stop = stop

    await make_worker(tmp_path, repository)._loop(QueueName.CONTROL, stop)

    assert repository.completed == [missing_job_task.id]
    assert repository.retried == []


async def test_jobs_queue_waits_for_provider_readiness(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    repository = FakeRepository()
    queued_task = task()
    repository.tasks = [queued_task]
    opencode = FakeOpenCode([False, True])
    stop = asyncio.Event()
    waits: list[asyncio.Event] = []
    value = make_worker(tmp_path, repository, opencode=opencode)

    async def wait(current_stop: asyncio.Event) -> None:
        waits.append(current_stop)

    async def handle(_current: Task) -> None:
        stop.set()

    monkeypatch.setattr(value, "_wait", wait)
    monkeypatch.setattr(value, "_handle", handle)

    await value._loop(QueueName.JOBS, stop)

    assert opencode.ready_calls == 2
    assert waits == [stop]
    assert repository.claimed_queues == [QueueName.JOBS]
    assert repository.completed == [queued_task.id]


async def test_run_recovers_before_starting_configured_consumers(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    value = make_worker(tmp_path, FakeRepository(), max_concurrent_jobs=3)
    calls: list[str | QueueName] = []

    async def recover() -> None:
        calls.append("recover")

    async def loop(queue: QueueName, _stop: asyncio.Event) -> None:
        calls.append(queue)

    monkeypatch.setattr(value, "_recover", recover)
    monkeypatch.setattr(value, "_loop", loop)

    await value.run(asyncio.Event())

    assert calls[0] == "recover"
    assert calls.count(QueueName.CONTROL) == 1
    assert calls.count(QueueName.JOBS) == 3


@pytest.mark.parametrize(
    ("permitted", "event_type", "suffix"),
    [
        (True, PermissionGranted, "permission-granted"),
        (False, PermissionDenied, "permission-denied"),
    ],
)
async def test_authorize_records_permission_outcome(
    tmp_path: Path,
    permitted: bool,
    event_type: type[PermissionGranted] | type[PermissionDenied],
    suffix: str,
) -> None:
    current = job()
    repository = FakeRepository(current)
    gitea = FakeGitea(permitted=permitted)
    authorize = task(TaskKind.AUTHORIZE, QueueName.CONTROL)

    await make_worker(tmp_path, repository, gitea=gitea)._authorize(authorize, current)

    assert gitea.permission_calls == [("org", "repo", "alice")]
    assert repository.events[0][0] == f"task:{authorize.id}:{suffix}"
    assert isinstance(repository.events[0][1], event_type)


async def test_authorize_ignores_job_after_received_status(tmp_path: Path) -> None:
    current = job(status=JobStatus.QUEUED)
    repository = FakeRepository(current)
    gitea = FakeGitea()

    await make_worker(tmp_path, repository, gitea=gitea)._authorize(
        task(TaskKind.AUTHORIZE, QueueName.CONTROL), current
    )

    assert gitea.permission_calls == []
    assert repository.events == []


async def test_authorize_integrates_with_real_repository(
    tmp_path: Path,
) -> None:
    repository = Repository(tmp_path / "state.sqlite3", MIGRATIONS)
    await repository.initialize()
    accepted = await repository.accept(
        IncomingCommand(
            delivery_id="delivery-real",
            comment_id=1,
            repo_owner="org",
            repo_name="repo",
            issue_number=1,
            pr_number=None,
            requester="alice",
            body="/agent plan write tests",
        )
    )
    authorize = await repository.claim_task(QueueName.CONTROL)
    assert authorize is not None
    assert authorize.kind is TaskKind.AUTHORIZE

    await make_worker(tmp_path, repository)._handle(authorize)

    persisted = await repository.get_job(accepted.job.id)
    assert persisted is not None
    assert persisted.status is JobStatus.QUEUED
    assert persisted.kind is JobKind.PLAN
    assert persisted.message == "write tests"
    execute = await repository.claim_task(QueueName.JOBS)
    assert execute is not None
    assert execute.kind is TaskKind.EXECUTE


async def test_execute_completes_with_workflow_comment(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    current = job(status=JobStatus.QUEUED)
    repository = FakeRepository(current)
    value = make_worker(tmp_path, repository)

    async def dispatch(_job: Job, _run: object, _services: object) -> str:
        return "final workflow body"

    monkeypatch.setattr("agentci.worker.dispatch", dispatch)
    execute = task()

    await value._execute(execute, current)

    assert [event_id for event_id, _ in repository.events] == [
        f"task:{execute.id}:started",
        f"task:{execute.id}:completed",
    ]
    assert isinstance(repository.events[0][1], JobStarted)
    completed = repository.events[1][1]
    assert isinstance(completed, JobCompleted)
    assert completed.comment_body == "final workflow body"


async def test_execute_records_expected_rejection(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    current = job(status=JobStatus.QUEUED)
    repository = FakeRepository(current)

    async def dispatch(_job: Job, _run: object, _services: object) -> str:
        raise JobRejected("pull request is closed")

    monkeypatch.setattr("agentci.worker.dispatch", dispatch)

    await make_worker(tmp_path, repository)._execute(task(), current)

    rejected = repository.events[-1][1]
    assert isinstance(rejected, RejectedEvent)
    assert rejected.reason == "pull request is closed"


async def test_execute_records_failure_at_latest_persisted_stage(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    current = job(status=JobStatus.QUEUED)
    repository = FakeRepository(current)

    async def dispatch(_job: Job, _run: object, _services: object) -> str:
        assert repository.job is not None
        repository.job = replace(repository.job, stage="cloning")
        raise RuntimeError("provider\nfailed")

    monkeypatch.setattr("agentci.worker.dispatch", dispatch)

    await make_worker(tmp_path, repository)._execute(task(), current)

    failed = repository.events[-1][1]
    assert isinstance(failed, JobFailed)
    assert failed.stage == "cloning"
    assert failed.error == "RuntimeError: provider failed"


async def test_execute_marks_running_job_interrupted_after_restart(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    current = job(status=JobStatus.RUNNING)
    repository = FakeRepository(current)

    async def unexpected_dispatch(_job: Job, _run: object, _services: object) -> str:
        pytest.fail("running jobs must not be dispatched again")

    monkeypatch.setattr("agentci.worker.dispatch", unexpected_dispatch)
    execute = task()

    await make_worker(tmp_path, repository)._execute(execute, current)

    assert repository.events[0][0] == f"task:{execute.id}:interrupted"
    assert isinstance(repository.events[0][1], ServiceRestarted)


@pytest.mark.parametrize("status", [JobStatus.RECEIVED, JobStatus.SUCCEEDED])
async def test_execute_ignores_ineligible_status(tmp_path: Path, status: JobStatus) -> None:
    current = job(status=status)
    repository = FakeRepository(current)

    await make_worker(tmp_path, repository)._execute(task(), current)

    assert repository.events == []


async def test_reconcile_updates_linked_comment_in_place(tmp_path: Path) -> None:
    current = job(status=JobStatus.QUEUED, accepted_comment_id=19)
    repository = FakeRepository(current)
    gitea = FakeGitea()

    await make_worker(tmp_path, repository, gitea=gitea)._reconcile(
        task(TaskKind.RECONCILE_COMMENT, QueueName.CONTROL), current
    )

    assert gitea.update_calls == [("org", "repo", 19, render_job_comment(current))]
    assert gitea.issue_comment_calls == []
    assert gitea.create_calls == []
    assert repository.events == []


async def test_reconcile_finds_oldest_matching_bot_comment_and_links_it(
    tmp_path: Path,
) -> None:
    current = job(status=JobStatus.QUEUED, accepted_comment_id=99)
    marker = f"<!-- agentci:job id={current.id} -->"
    repository = FakeRepository(current)
    gitea = FakeGitea()
    gitea.update_results = [False, True]
    gitea.comments = [
        CommentInfo(8, "agentci", f"{marker}\nnewer", ""),
        CommentInfo(3, "AgentCI", f"{marker}\nolder", ""),
        CommentInfo(1, "alice", f"{marker}\nnot the bot", ""),
        CommentInfo(2, "agentci", f"prefix {marker}", ""),
    ]
    reconcile = task(TaskKind.RECONCILE_COMMENT, QueueName.CONTROL)

    await make_worker(tmp_path, repository, gitea=gitea)._reconcile(reconcile, current)

    assert gitea.issue_comment_calls == [("org", "repo", 1)]
    assert gitea.create_calls == []
    assert [call[2] for call in gitea.update_calls] == [99, 3]
    assert repository.events[0][0] == f"task:{reconcile.id}:comment:3"
    linked = repository.events[0][1]
    assert isinstance(linked, CommentLinked)
    assert linked.comment_id == 3


async def test_reconcile_creates_links_and_populates_missing_comment(tmp_path: Path) -> None:
    current = job(status=JobStatus.SUCCEEDED)
    repository = FakeRepository(current)
    gitea = FakeGitea()
    gitea.comments = [CommentInfo(1, "alice", "unrelated", "")]
    reconcile = task(TaskKind.RECONCILE_COMMENT, QueueName.CONTROL)

    await make_worker(tmp_path, repository, gitea=gitea)._reconcile(reconcile, current)

    body = render_job_comment(current)
    assert gitea.create_calls == [("org", "repo", 1, body)]
    assert gitea.update_calls == [("org", "repo", 42, body)]
    assert repository.events[0][0] == f"task:{reconcile.id}:comment:42"
    linked = repository.events[0][1]
    assert isinstance(linked, CommentLinked)
    assert linked.comment_id == 42


async def test_recovery_resets_tasks_before_interrupting_running_jobs(tmp_path: Path) -> None:
    repository = FakeRepository()
    repository.running = [
        job(job_id="first", status=JobStatus.RUNNING),
        job(job_id="second", status=JobStatus.RUNNING),
    ]

    await make_worker(tmp_path, repository)._recover()

    assert repository.operations == [
        "recover_tasks",
        "running_jobs",
        "apply:recovery:first:service-restarted",
        "apply:recovery:second:service-restarted",
    ]
    assert all(isinstance(event, ServiceRestarted) for _, event in repository.events)


@pytest.mark.parametrize(
    ("primary", "reviewer", "expected"),
    [
        ("primary", "reviewer", {"primary", "reviewer"}),
        ("same-session", "same-session", {"same-session"}),
    ],
)
async def test_abort_collects_and_deduplicates_workflow_sessions(
    tmp_path: Path,
    primary: str,
    reviewer: str,
    expected: set[str],
) -> None:
    workspace = tmp_path / "workflow" / "repo"
    workflow = Workflow(
        id="flow",
        kind=WorkflowKind.IMPLEMENT,
        repo_owner="org",
        repo_name="repo",
        issue_number=1,
        workspace_path=workspace,
        base_sha="base",
        primary_session_id=primary,
        reviewer_session_id=reviewer,
    )
    opencode = FakeOpenCode()

    await make_worker(
        tmp_path, FakeRepository(workflow=workflow), opencode=opencode
    )._abort_job_sessions(job(workflow_id="flow"))

    assert set(opencode.aborted) == {(session, workspace) for session in expected}


async def test_abort_uses_one_shot_workspace_without_workflow(tmp_path: Path) -> None:
    opencode = FakeOpenCode()

    await make_worker(tmp_path, FakeRepository(), opencode=opencode)._abort_job_sessions(
        job(workflow_id="missing", session_id="session")
    )

    assert opencode.aborted == [("session", tmp_path / "fix-job" / "repo")]


async def test_abort_does_not_mix_one_shot_session_into_existing_workflow(
    tmp_path: Path,
) -> None:
    workflow = Workflow(
        id="flow",
        kind=WorkflowKind.IMPLEMENT,
        repo_owner="org",
        repo_name="repo",
        issue_number=1,
        workspace_path=tmp_path / "workflow" / "repo",
        base_sha="base",
    )
    opencode = FakeOpenCode()

    await make_worker(
        tmp_path, FakeRepository(workflow=workflow), opencode=opencode
    )._abort_job_sessions(job(workflow_id="flow", session_id="one-shot"))

    assert opencode.aborted == []


async def test_abort_without_persisted_sessions_is_noop(tmp_path: Path) -> None:
    opencode = FakeOpenCode()

    await make_worker(tmp_path, FakeRepository(), opencode=opencode)._abort_job_sessions(job())

    assert opencode.aborted == []


def test_safe_error_is_single_line_and_bounded() -> None:
    value = _safe_error(RuntimeError("bad\n" + "x" * 2000))

    assert "\n" not in value
    assert len(value) == 1000