rein-aharness/tests/test_native_limits.py

from __future__ import annotations
import json
import subprocess
from unittest.mock import MagicMock, patch
import pytest
from llm_connect.models import BudgetTracker, RunConfig
from rein_aharness.adapter import AgenticClaudeCodeAdapter
from rein_aharness.native_limits import NativeLimitError, terminal_accounting
from rein_aharness.taskspec import TaskSpec, TaskSpecError
from test_adapter import _fake_proc


def terminal(**overrides):
    return {
        "type": "result",
        "subtype": "success",
        "is_error": False,
        "result": "done",
        "num_turns": 2,
        "total_cost_usd": 0.1,
        "usage": {"input_tokens": 10, "output_tokens": 2, "cache_read_input_tokens": 3},
        **overrides,
    }


@pytest.mark.parametrize("stream", [False, True])
def test_native_controls_reach_cli_and_account_usage(tmp_path, stream):
    adapter = AgenticClaudeCodeAdapter(
        workdir=tmp_path, on_tool_event=(lambda e: None) if stream else None
    )
    config = RunConfig(
        timeout_seconds=30,
        model_params={"max_budget_usd": 0.25, "max_turns": 3},
        budget_tracker=BudgetTracker(total=100),
    )
    payload = json.dumps(terminal())
    proc = _fake_proc([payload + "\n"]) if stream else MagicMock()
    if not stream:
        proc.communicate.return_value = (payload, "")
        proc.returncode = 0
    with (
        patch(
            "rein_aharness.adapter.subprocess.run",
            return_value=subprocess.CompletedProcess(
                [], 0, "2.1.266 (Claude Code)\n", ""
            ),
        ),
        patch("rein_aharness.adapter.subprocess.Popen", return_value=proc) as invoke,
    ):
        result = adapter.execute_prompt("task", config)
    argv = invoke.call_args.args[0]
    assert argv[argv.index("--max-budget-usd") + 1] == "0.25"
    assert argv[argv.index("--max-turns") + 1] == "3"
    assert result.usage["total_tokens"] == config.budget_tracker.spent == 15
    assert result.metadata["cost_usd"] == 0.1


@pytest.mark.parametrize(
    "changes",
    [
        {"total_cost_usd": None},
        {"total_cost_usd": float("nan")},
        {"total_cost_usd": True},
        {"total_cost_usd": 0.3},
        {"num_turns": 4},
        {"usage": {}},
        {"subtype": "error_max_budget_usd", "is_error": True},
        {"subtype": "error_max_turns", "is_error": True},
    ],
)
def test_unaccounted_or_exhausted_result_cannot_succeed(changes):
    with pytest.raises(NativeLimitError):
        terminal_accounting(terminal(**changes), max_budget_usd=0.25, max_turns=3)


def test_error_preserves_only_bounded_cost():
    with pytest.raises(NativeLimitError) as caught:
        terminal_accounting(
            terminal(subtype="error_max_budget_usd", is_error=True, errors=["secret"]),
            max_budget_usd=0.25,
            max_turns=3,
        )
    assert caught.value.cost_usd == 0.1 and "secret" not in str(caught.value)


@pytest.mark.parametrize("version", ["2.1.216 (Claude Code)", "2.1.265 (Claude Code)", "unrecognized"])
def test_old_cli_refuses_before_prompt(tmp_path, version):
    adapter = AgenticClaudeCodeAdapter(workdir=tmp_path)
    with (
        patch(
            "rein_aharness.adapter.subprocess.run",
            return_value=subprocess.CompletedProcess([], 0, version, ""),
        ),
        patch("rein_aharness.adapter.subprocess.Popen") as invoke,
    ):
        with pytest.raises(NativeLimitError, match="2.1.266"):
            adapter.execute_prompt(
                "task", RunConfig(model_params={"max_budget_usd": 0.25})
            )
    invoke.assert_not_called()


@pytest.mark.parametrize(
    "limits",
    [
        {"max_budget_usd": True},
        {"max_budget_usd": -1},
        {"max_budget_usd": float("inf")},
        {"max_turns": True},
        {"max_turns": 1.5},
    ],
)
def test_invalid_task_limits_refuse(tmp_path, limits):
    with pytest.raises(TaskSpecError):
        TaskSpec(title="task", description="d", target_repo=tmp_path, **limits)


def test_runner_forwards_task_limits(tmp_path):
    from test_runner import CommittingAdapter, _make_repo
    from rein_aharness.runner import run_task

    repo = _make_repo(tmp_path)
    adapter = CommittingAdapter(repo)
    result = run_task(
        TaskSpec(
            title="task",
            description="d",
            target_repo=repo,
            max_budget_usd=0.25,
            max_turns=3,
        ),
        adapter=adapter,
        report_to_hub=False,
        write_metrics=False,
    )
    assert result.ok
    assert adapter.configs[0].model_params == {"max_budget_usd": 0.25, "max_turns": 3}


@pytest.mark.parametrize("stream", [False, True])
def test_controlled_session_uses_closed_tools_and_skips_ambient_config(tmp_path, stream):
    adapter = AgenticClaudeCodeAdapter(workdir=tmp_path, on_tool_event=(lambda e: None) if stream else None)
    cmd = adapter._build_command(RunConfig(model_params={"max_budget_usd": 1, "max_turns": 4}))
    assert cmd[cmd.index("--permission-mode")+1] == "dontAsk"
    assert cmd[cmd.index("--tools")+1] == "Read,Write,Edit,Glob,Grep,Bash"
    assert cmd[cmd.index("--setting-sources")+1] == ""
    assert json.loads(cmd[cmd.index("--mcp-config")+1]) == {"mcpServers": {}}
    assert "--bare" in cmd and "--strict-mcp-config" in cmd
    assert cmd[cmd.index("--disallowedTools")+1] == "mcp__*"


def test_invalid_registered_tool_rule_refuses_before_prompt(tmp_path):
    from rein_aharness.profiles import ToolProfile
    adapter = AgenticClaudeCodeAdapter(workdir=tmp_path, tool_profile=ToolProfile("invalid", "test", "Read,*", "green"))
    with patch("rein_aharness.adapter.subprocess.Popen") as invoke:
        with pytest.raises(ValueError, match="builtin rule"):
            adapter.execute_prompt("task", RunConfig(model_params={"max_turns": 1}))
    invoke.assert_not_called()