rein-aharness/tests/test_runner.py

from __future__ import annotations

import json
import subprocess
from pathlib import Path

import pytest
import yaml

from rein_aharness.manifest import HARNESS_MAJOR
from rein_aharness.runner import RunResult, run_task
from rein_aharness.taskspec import TaskSpec, TaskSpecError


def _make_repo(tmp_path: Path) -> Path:
    repo = tmp_path / "sandbox"
    repo.mkdir()
    subprocess.run(["git", "init", "-q"], cwd=repo, check=True)
    (repo / "README.md").write_text("sandbox\n")
    subprocess.run(["git", "add", "."], cwd=repo, check=True)
    subprocess.run(
        ["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", "init"],
        cwd=repo,
        check=True,
    )
    return repo


def _write_manifest(repo: Path, agents: dict) -> None:
    kaizen = repo / ".kaizen"
    kaizen.mkdir(exist_ok=True)
    data = {"version": "1", "harness": HARNESS_MAJOR, "agents": agents}
    (kaizen / "schedule.yml").write_text(
        yaml.safe_dump(data, sort_keys=False), encoding="utf-8"
    )


class CommittingAdapter:
    """Fake adapter that simulates a session which commits."""

    def __init__(self, repo: Path):
        self.repo = repo
        self.prompts: list[str] = []
        self.configs: list = []

    def execute_prompt(self, prompt, config):
        self.prompts.append(prompt)
        self.configs.append(config)
        (self.repo / "HELLO.md").write_text("hello\n")
        subprocess.run(["git", "add", "."], cwd=self.repo, check=True)
        subprocess.run(
            ["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", "task done"],
            cwd=self.repo,
            check=True,
        )
        from llm_connect.models import LLMResponse

        return LLMResponse(content="done", model="fake", usage={}, finish_reason="stop")


class IdleAdapter:
    def execute_prompt(self, prompt, config):
        from llm_connect.models import LLMResponse

        return LLMResponse(content="nothing to do", model="fake", usage={}, finish_reason="stop")


class BudgetBlowingAdapter:
    def execute_prompt(self, prompt, config):
        from llm_connect.exceptions import LLMBudgetExceededError

        if config.budget_tracker is not None:
            # Simulate preflight/exhaustion the way adapters do.
            config.budget_tracker.consume(config.budget_tracker.total)
            config.budget_tracker.consume(1)
        raise LLMBudgetExceededError(
            "Token budget exceeded",
            total=1,
            spent=1,
            requested=1,
        )


def _spec(repo: Path, agent: str = "coach") -> TaskSpec:
    return TaskSpec(
        title="write hello",
        description="create HELLO.md",
        target_repo=repo,
        agent=agent,
    )


def test_run_task_success_when_session_commits(tmp_path) -> None:
    repo = _make_repo(tmp_path)
    adapter = CommittingAdapter(repo)

    result = run_task(
        _spec(repo), adapter=adapter, report_to_hub=False, write_metrics=True
    )

    assert isinstance(result, RunResult)
    assert result.ok is True
    assert result.committed is True
    assert result.head_before != result.head_after
    assert result.tool_profile == "green-commit-only"
    assert "write hello" in adapter.prompts[0]
    assert "Never push" in adapter.prompts[0]
    assert "Tool profile for this run: green-commit-only" in adapter.prompts[0]

    metrics_path = repo / ".kaizen" / "metrics" / "coach" / "executions.jsonl"
    assert metrics_path.is_file()
    rec = json.loads(metrics_path.read_text().strip().splitlines()[-1])
    assert rec["success"] is True
    assert rec["committed"] is True


def test_run_task_fails_without_commit(tmp_path) -> None:
    repo = _make_repo(tmp_path)

    result = run_task(
        _spec(repo), adapter=IdleAdapter(), report_to_hub=False, write_metrics=False
    )

    assert result.ok is False
    assert result.committed is False
    assert result.reason == "session completed without committing"


def test_run_task_refuses_unknown_tool_profile(tmp_path) -> None:
    repo = _make_repo(tmp_path)
    _write_manifest(
        repo,
        {
            "coach": {
                "cadence": "daily",
                "tool_profile": "not-a-real-profile",
            }
        },
    )

    result = run_task(
        _spec(repo), adapter=IdleAdapter(), report_to_hub=False, write_metrics=False
    )

    assert result.ok is False
    assert result.reason.startswith("refused:")
    assert "not-a-real-profile" in result.reason


def test_run_task_resolves_manifest_profile_and_budget(tmp_path) -> None:
    repo = _make_repo(tmp_path)
    _write_manifest(
        repo,
        {
            "coach": {
                "cadence": "daily",
                "lane": "blue",
                "tool_profile": "blue-mail-triage",
                "budget": 25000,
            }
        },
    )
    adapter = CommittingAdapter(repo)

    result = run_task(
        _spec(repo), adapter=adapter, report_to_hub=False, write_metrics=False
    )

    assert result.ok is True
    assert result.tool_profile == "blue-mail-triage"
    assert result.budget_tokens == 25000
    assert adapter.configs[0].budget_tracker is not None
    assert adapter.configs[0].budget_tracker.total == 25000
    assert "blue-mail-triage" in adapter.prompts[0]


def test_run_task_budget_exhaustion_fails(tmp_path) -> None:
    repo = _make_repo(tmp_path)
    _write_manifest(
        repo,
        {
            "coach": {
                "cadence": "daily",
                "tool_profile": "green-commit-only",
                "budget": 10,
            }
        },
    )

    result = run_task(
        _spec(repo),
        adapter=BudgetBlowingAdapter(),
        report_to_hub=False,
        write_metrics=False,
    )

    assert result.ok is False
    assert "session failed" in result.reason
    assert result.budget_tokens == 10


def test_taskspec_rejects_non_repo(tmp_path) -> None:
    spec_file = tmp_path / "task.json"
    spec_file.write_text(
        '{"title": "x", "description": "y", "target_repo": "%s"}' % tmp_path
    )
    with pytest.raises(TaskSpecError, match="not a git repository"):
        TaskSpec.from_file(spec_file)


def test_run_task_wires_on_tool_event_when_emit_tool_events(tmp_path, monkeypatch) -> None:
    repo = _make_repo(tmp_path)
    spec = _spec(repo)
    captured: dict = {}

    class FakeStreamingAdapter:
        def __init__(self, workdir, tool_profile, on_tool_event=None):
            captured["on_tool_event"] = on_tool_event

        def execute_prompt(self, prompt, config):
            captured["on_tool_event"](
                {"type": "assistant", "message": {"content": [{"type": "tool_use", "name": "Read"}]}}
            )
            (repo / "HELLO.md").write_text("hi\n")
            subprocess.run(["git", "add", "."], cwd=repo, check=True)
            subprocess.run(
                ["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", "task"],
                cwd=repo,
                check=True,
            )
            from llm_connect.models import LLMResponse

            return LLMResponse(content="done", model="m", usage={}, finish_reason="stop", metadata={})

    monkeypatch.setattr("rein_aharness.adapter.AgenticClaudeCodeAdapter", FakeStreamingAdapter)

    result = run_task(spec, report_to_hub=False, write_metrics=False, emit_tool_events=True)

    assert captured["on_tool_event"] is not None
    assert len(result.tool_events) == 1
    assert result.tool_events[0]["message"]["content"][0]["name"] == "Read"


def test_run_task_default_adapter_gets_no_callback_when_not_requested(tmp_path, monkeypatch) -> None:
    repo = _make_repo(tmp_path)
    spec = _spec(repo)
    captured: dict = {}

    class FakeAdapter:
        def __init__(self, workdir, tool_profile, on_tool_event=None):
            captured["on_tool_event"] = on_tool_event

        def execute_prompt(self, prompt, config):
            (repo / "HELLO.md").write_text("hi\n")
            subprocess.run(["git", "add", "."], cwd=repo, check=True)
            subprocess.run(
                ["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", "task"],
                cwd=repo,
                check=True,
            )
            from llm_connect.models import LLMResponse

            return LLMResponse(content="done", model="m", usage={}, finish_reason="stop", metadata={})

    monkeypatch.setattr("rein_aharness.adapter.AgenticClaudeCodeAdapter", FakeAdapter)

    result = run_task(spec, report_to_hub=False, write_metrics=False)

    assert captured["on_tool_event"] is None
    assert result.tool_events == []