agentci/tests/test_code_review.py

from pathlib import Path
from types import SimpleNamespace
from typing import Any, cast

from pydantic import BaseModel

from agentci.engine.model import Workflow, WorkflowKind
from agentci.engine.run import JobRun
from agentci.workflows.model import (
    AgentResult,
    PlanArtifact,
    ReviewFinding,
    ReviewReport,
    ReviewSeverity,
)
from agentci.workflows.review import (
    review_implementation_loop,
    review_implementation_once,
    review_plan_loop,
    review_plan_once,
)
from agentci.workflows.services import WorkflowServices


def serious_report(summary: str = "Needs work") -> ReviewReport:
    return ReviewReport(
        summary=summary,
        findings=[
            ReviewFinding(
                severity=ReviewSeverity.MAJOR,
                title="Missing check",
                detail="A check is absent.",
                recommendation="Add it.",
            )
        ],
    )


def clean_report() -> ReviewReport:
    return ReviewReport(summary="Ready", findings=[])


def minor_report() -> ReviewReport:
    return ReviewReport(
        summary="Optional improvement",
        findings=[
            ReviewFinding(
                severity=ReviewSeverity.MINOR,
                title="Clarify wording",
                detail="The wording could be clearer.",
                recommendation="Tighten it when convenient.",
            )
        ],
    )


class RecordingOpenCode:
    def __init__(self, responses: list[BaseModel]) -> None:
        self.responses = list(responses)
        self.created_sessions: list[tuple[Path, str]] = []
        self.resume_calls: list[dict[str, Any]] = []

    async def create_session(self, workspace: Path, title: str) -> str:
        self.created_sessions.append((workspace, title))
        return f"{title}-session"

    async def resume(self, **values: Any) -> BaseModel:
        self.resume_calls.append(values)
        response = self.responses.pop(0)
        assert isinstance(response, values["result_type"])
        return response


class RecordingRepository:
    def __init__(self) -> None:
        self.saved_workflows: list[Workflow] = []

    async def save_workflow(self, workflow: Workflow) -> None:
        self.saved_workflows.append(workflow)


class RecordingRun(JobRun):
    def __init__(self) -> None:
        self.stages: list[str] = []

    async def stage(self, stage: str) -> None:
        self.stages.append(stage)


class RecordingPrompts:
    def __init__(self) -> None:
        self.calls: list[tuple[str, dict[str, str]]] = []

    def render(self, name: str, **values: str) -> str:
        self.calls.append((name, values))
        return f"rendered {name}"


def workflow(*, reviewer_session_id: str | None = None) -> Workflow:
    return Workflow(
        id="flow",
        kind=WorkflowKind.IMPLEMENT,
        repo_owner="org",
        repo_name="repo",
        issue_number=1,
        workspace_path=Path("/workspace/repo"),
        base_sha="abc",
        primary_session_id="primary-session",
        reviewer_session_id=reviewer_session_id,
    )


def objects(
    responses: list[BaseModel],
    *,
    plan_rounds: int = 4,
    implementation_rounds: int = 3,
) -> tuple[
    RecordingOpenCode,
    RecordingRepository,
    RecordingPrompts,
    WorkflowServices,
]:
    opencode = RecordingOpenCode(responses)
    repository = RecordingRepository()
    prompts = RecordingPrompts()
    services = cast(
        WorkflowServices,
        SimpleNamespace(
            settings=SimpleNamespace(
                plan_review_rounds=plan_rounds,
                plan_model="provider/plan",
                plan_variant="high",
                implement_review_rounds=implementation_rounds,
                implement_model="provider/implement",
                implement_variant="high",
            ),
            opencode=opencode,
            repository=repository,
            prompts=prompts,
            development=SimpleNamespace(description="Python 3.13"),
        ),
    )
    return opencode, repository, prompts, services


async def test_implementation_loop_persists_reviewed_revision_and_stops_clean() -> None:
    revised = AgentResult(summary_markdown="revision 1", tests=["pytest: passed"])
    opencode, repository, prompts, services = objects(
        [serious_report(), revised, clean_report()], implementation_rounds=4
    )
    run = RecordingRun()
    original = workflow()

    updated, result, report = await review_implementation_loop(
        original,
        "issue context",
        "canonical plan",
        AgentResult(summary_markdown="initial", tests=[]),
        run,
        services,
    )

    assert result == revised
    assert report == clean_report()
    assert original.reviewer_session_id is None
    assert updated.reviewer_session_id == "implementation-review-session"
    assert updated.artifact == revised.model_dump_json()
    assert updated.review_json == clean_report().model_dump_json()
    assert repository.saved_workflows[-1] == updated
    assert repository.saved_workflows[0].reviewer_session_id == ("implementation-review-session")
    assert run.stages == [
        "reviewing implementation 1/4",
        "reviewing implementation 2/4",
    ]
    assert opencode.created_sessions == [(original.workspace_path, "implementation-review")]
    assert [call["session_id"] for call in opencode.resume_calls] == [
        "implementation-review-session",
        "primary-session",
        "implementation-review-session",
    ]
    assert [call["result_type"] for call in opencode.resume_calls] == [
        ReviewReport,
        AgentResult,
        ReviewReport,
    ]
    assert [name for name, _ in prompts.calls] == [
        "implementation_review",
        "implementation_revision",
        "implementation_review",
    ]
    assert opencode.responses == []


async def test_implementation_loop_never_makes_unreviewed_final_revision() -> None:
    final_report = serious_report("Still failing after the last review")
    opencode, repository, _, services = objects(
        [
            serious_report("round 1"),
            AgentResult(summary_markdown="revision 1", tests=[]),
            serious_report("round 2"),
            AgentResult(summary_markdown="revision 2", tests=[]),
            final_report,
        ],
        implementation_rounds=3,
    )
    run = RecordingRun()

    updated, result, report = await review_implementation_loop(
        workflow(),
        "issue context",
        "canonical plan",
        AgentResult(summary_markdown="initial", tests=[]),
        run,
        services,
    )

    assert result.summary_markdown == "revision 2"
    assert report is final_report
    assert updated.artifact == result.model_dump_json()
    assert updated.review_json == final_report.model_dump_json()
    assert repository.saved_workflows[-1] == updated
    assert [call["result_type"] for call in opencode.resume_calls].count(ReviewReport) == 3
    assert [call["result_type"] for call in opencode.resume_calls].count(AgentResult) == 2
    assert run.stages == [
        "reviewing implementation 1/3",
        "reviewing implementation 2/3",
        "reviewing implementation 3/3",
    ]
    assert opencode.responses == []


async def test_implementation_review_once_reuses_existing_reviewer_session() -> None:
    opencode, repository, prompts, services = objects([clean_report()])
    existing = workflow(reviewer_session_id="existing-reviewer")

    updated, report = await review_implementation_once(
        existing,
        issue_context="issue context",
        plan="canonical plan",
        pull_context="pull request context",
        services=services,
    )

    assert updated is existing
    assert report == clean_report()
    assert opencode.created_sessions == []
    assert repository.saved_workflows == []
    assert opencode.resume_calls == [
        {
            "session_id": "existing-reviewer",
            "prompt": "rendered implementation_review",
            "model": "provider/implement",
            "variant": "high",
            "workspace": existing.workspace_path,
            "schema_name": "review.json",
            "result_type": ReviewReport,
        }
    ]
    assert prompts.calls == [
        (
            "implementation_review",
            {
                "issue_context": "issue context",
                "artifact": "canonical plan",
                "pull_context": "pull request context",
            },
        )
    ]


async def test_plan_loop_revises_serious_finding_then_persists_clean_result() -> None:
    revised = PlanArtifact(plan_markdown="Revised plan")
    opencode, repository, prompts, services = objects(
        [serious_report(), revised, clean_report()], plan_rounds=4
    )
    run = RecordingRun()
    original = workflow()
    initial = PlanArtifact(plan_markdown="Initial plan")

    updated, artifact, report = await review_plan_loop(
        original, "issue context", initial, run, services
    )

    assert artifact is revised
    assert report == clean_report()
    assert updated.reviewer_session_id == "plan-review-session"
    assert updated.artifact == "Revised plan"
    assert updated.review_json == clean_report().model_dump_json()
    assert repository.saved_workflows[-1] == updated
    assert run.stages == ["reviewing plan 1/4", "reviewing plan 2/4"]
    assert [call["session_id"] for call in opencode.resume_calls] == [
        "plan-review-session",
        "primary-session",
        "plan-review-session",
    ]
    assert [name for name, _ in prompts.calls] == [
        "plan_review",
        "plan_revision",
        "plan_review",
    ]
    assert opencode.responses == []


async def test_plan_loop_stops_at_round_boundary_without_unreviewed_revision() -> None:
    final_report = serious_report("round 2")
    opencode, repository, _, services = objects(
        [
            serious_report("round 1"),
            PlanArtifact(plan_markdown="Only revision"),
            final_report,
        ],
        plan_rounds=2,
    )
    run = RecordingRun()

    updated, artifact, report = await review_plan_loop(
        workflow(),
        "issue context",
        PlanArtifact(plan_markdown="Initial plan"),
        run,
        services,
    )

    assert artifact.plan_markdown == "Only revision"
    assert report is final_report
    assert updated.artifact == "Only revision"
    assert updated.review_json == final_report.model_dump_json()
    assert repository.saved_workflows[-1] == updated
    assert [call["result_type"] for call in opencode.resume_calls] == [
        ReviewReport,
        PlanArtifact,
        ReviewReport,
    ]
    assert run.stages == ["reviewing plan 1/2", "reviewing plan 2/2"]
    assert opencode.responses == []


async def test_plan_review_once_creates_and_persists_reviewer_session() -> None:
    opencode, repository, prompts, services = objects([clean_report()])
    original = workflow()

    updated, report = await review_plan_once(
        original,
        "issue context",
        PlanArtifact(plan_markdown="Plan body"),
        services,
    )

    assert report == clean_report()
    assert updated is not original
    assert original.reviewer_session_id is None
    assert updated.reviewer_session_id == "plan-review-session"
    assert repository.saved_workflows == [updated]
    assert opencode.created_sessions == [(original.workspace_path, "plan-review")]
    assert opencode.resume_calls[0]["session_id"] == "plan-review-session"
    assert opencode.resume_calls[0]["result_type"] is ReviewReport
    assert prompts.calls == [
        (
            "plan_review",
            {"context": "issue context", "artifact": "Plan body"},
        )
    ]


async def test_minor_findings_end_review_loop_without_revision() -> None:
    opencode, repository, _, services = objects([minor_report()], implementation_rounds=5)
    run = RecordingRun()
    initial = AgentResult(summary_markdown="initial", tests=[])

    updated, result, report = await review_implementation_loop(
        workflow(),
        "issue context",
        "canonical plan",
        initial,
        run,
        services,
    )

    assert result is initial
    assert report == minor_report()
    assert not report.has_serious_findings
    assert updated.artifact == initial.model_dump_json()
    assert updated.review_json == minor_report().model_dump_json()
    assert repository.saved_workflows[-1] == updated
    assert [call["result_type"] for call in opencode.resume_calls] == [ReviewReport]
    assert run.stages == ["reviewing implementation 1/5"]