tests/test_pilot.py

from __future__ import annotations

import asyncio
import base64
import hashlib
import json

import pytest

import autocad_bench.harness.run as pilot_module
from autocad_bench.harness.agent_loop import ModelTurn
from autocad_bench.sandbox.broker.models import DesktopScreenshot
from autocad_bench.harness.run import _LiveTraceWriter, _run, build_parser

_PNG = base64.b64decode(
    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwC"
    "AAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="
)


def test_pilot_defaults_to_requested_openai_model() -> None:
    args = build_parser().parse_args(
        ["--broker-url", "http://127.0.0.1:8765", "--output-dir", "unused"]
    )
    assert args.provider == "openai"
    assert args.model_id == "gpt-5.6-sol"
    assert args.reasoning_effort == "xhigh"
    assert args.reasoning_max_tokens is None
    assert args.tool_choice_mode == "specified"
    assert args.invalid_action_limit == 10
    assert args.expected_broker_version == "windows-autocad-2019-v10"


def test_pilot_supports_disabling_only_the_rollout_wall_timeout() -> None:
    args = build_parser().parse_args(
        [
            "--broker-url",
            "http://127.0.0.1:8765",
            "--output-dir",
            "unused",
            "--no-wall-timeout",
        ]
    )
    assert args.no_wall_timeout is True
    assert args.wall_timeout_s is None


def test_pilot_missing_openai_key_fails_before_creating_run_directory(
    tmp_path, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
    output_dir = tmp_path / "must-not-exist"
    args = build_parser().parse_args(
        [
            "--broker-url",
            "http://127.0.0.1:8765",
            "--output-dir",
            str(output_dir),
        ]
    )
    with pytest.raises(SystemExit, match="OPENAI_API_KEY"):
        asyncio.run(_run(args))
    assert not output_dir.exists()


def test_pilot_supports_openrouter_anthropic_fireworks_and_mantle_providers() -> None:
    openrouter = build_parser().parse_args(
        [
            "--broker-url",
            "http://127.0.0.1:8765",
            "--output-dir",
            "unused",
            "--provider",
            "openrouter",
            "--model-id",
            "moonshotai/kimi-k3",
            "--reasoning-effort",
            "max",
            "--tool-choice-mode",
            "required",
        ]
    )
    anthropic = build_parser().parse_args(
        [
            "--broker-url",
            "http://127.0.0.1:8765",
            "--output-dir",
            "unused",
            "--provider",
            "anthropic",
            "--model-id",
            "claude-test",
        ]
    )
    mantle = build_parser().parse_args(
        [
            "--broker-url",
            "http://127.0.0.1:8765",
            "--output-dir",
            "unused",
            "--provider",
            "mantle",
            "--model-id",
            "moonshotai.kimi-k2.5",
        ]
    )
    fireworks = build_parser().parse_args(
        [
            "--broker-url",
            "http://127.0.0.1:8765",
            "--output-dir",
            "unused",
            "--provider",
            "fireworks",
            "--model-id",
            "accounts/fireworks/routers/kimi-k2p6-fast",
        ]
    )
    assert openrouter.provider == "openrouter"
    assert openrouter.openrouter_key_env == "OPENROUTER_API_KEY"
    assert openrouter.reasoning_effort == "max"
    assert openrouter.reasoning_max_tokens is None
    assert openrouter.tool_choice_mode == "required"
    assert anthropic.provider == "anthropic"
    assert anthropic.anthropic_key_env == "ANTHROPIC_API_KEY"
    assert mantle.provider == "mantle"
    assert mantle.aws_profile is None
    assert fireworks.provider == "fireworks"
    assert fireworks.fireworks_key_env == "FIREWORKS_API_KEY"


def test_pilot_missing_fireworks_key_fails_before_creating_run_directory(
    tmp_path, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.delenv("FIREWORKS_API_KEY", raising=False)
    output_dir = tmp_path / "must-not-exist"
    args = build_parser().parse_args(
        [
            "--broker-url",
            "http://127.0.0.1:8765",
            "--output-dir",
            str(output_dir),
            "--provider",
            "fireworks",
            "--model-id",
            "accounts/fireworks/routers/kimi-k2p6-fast",
        ]
    )
    with pytest.raises(SystemExit, match="FIREWORKS_API_KEY"):
        asyncio.run(_run(args))
    assert not output_dir.exists()


def test_pilot_missing_openrouter_key_fails_before_creating_run_directory(
    tmp_path, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
    output_dir = tmp_path / "must-not-exist"
    args = build_parser().parse_args(
        [
            "--broker-url",
            "http://127.0.0.1:8765",
            "--output-dir",
            str(output_dir),
            "--provider",
            "openrouter",
            "--model-id",
            "moonshotai/kimi-k3",
        ]
    )
    with pytest.raises(SystemExit, match="OPENROUTER_API_KEY"):
        asyncio.run(_run(args))
    assert not output_dir.exists()


def test_pilot_missing_anthropic_key_fails_before_creating_run_directory(
    tmp_path, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
    output_dir = tmp_path / "must-not-exist"
    args = build_parser().parse_args(
        [
            "--broker-url",
            "http://127.0.0.1:8765",
            "--output-dir",
            str(output_dir),
            "--provider",
            "anthropic",
            "--model-id",
            "claude-test",
        ]
    )
    with pytest.raises(SystemExit, match="ANTHROPIC_API_KEY"):
        asyncio.run(_run(args))
    assert not output_dir.exists()


def test_pilot_rejects_wrong_broker_before_creating_trace(
    tmp_path, monkeypatch: pytest.MonkeyPatch
) -> None:
    class WrongBroker:
        def __init__(self, *_: object, **__: object) -> None:
            pass

        async def health(self):
            return {
                "status": "ok",
                "broker_version": "windows-autocad-2019-v7",
                "input_driver": "pyautogui-0.9.54-v3",
                "input_driver_ready": True,
                "width": 1920,
                "height": 1080,
            }

    monkeypatch.setenv("OPENAI_API_KEY", "test-key")
    monkeypatch.setenv("AUTOCAD_BROKER_TOKEN", "test-token")
    monkeypatch.setattr(pilot_module, "HttpDesktopBroker", WrongBroker)
    output_dir = tmp_path / "must-not-exist"
    args = build_parser().parse_args(
        [
            "--broker-url",
            "http://127.0.0.1:8765",
            "--output-dir",
            str(output_dir),
        ]
    )

    with pytest.raises(SystemExit, match="rollout preflight"):
        asyncio.run(_run(args))

    assert not output_dir.exists()


def test_live_trace_writer_persists_frames_turns_receipts_and_atomic_state(
    tmp_path,
) -> None:
    output_dir = tmp_path / "run"
    output_dir.mkdir()
    screenshot = DesktopScreenshot(
        session_id="session",
        sequence=0,
        captured_at="2026-07-15T00:00:00Z",
        width=1920,
        height=1080,
        autocad_maximized=True,
        png_sha256=hashlib.sha256(_PNG).hexdigest(),
        png=_PNG,
    )
    writer = _LiveTraceWriter(
        output_dir=output_dir,
        task_id="task-001",
        attempt_id="pilot-test",
        provider="openai",
        model_id="gpt-5.6-sol",
        reasoning_effort="xhigh",
        reasoning_max_tokens=None,
        reference_png=_PNG,
    )
    writer("screenshot", screenshot)
    writer(
        "turn",
        ModelTurn(
            tool_use_id="call-1",
            action_payload={"action": "done"},
            stop_reason="completed",
            usage={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12},
            latency_ms=42.0,
            model_input={
                "model": "gpt-5.6-sol",
                "messages": [
                    {
                        "content": [
                            {"type": "text", "text": "very large repeated history"},
                            {
                                "type": "image",
                                "artifact_path": "frames/frame_0000.png",
                                "sha256": "abc",
                                "bytes": 123,
                            },
                        ],
                    }
                ],
            },
        ),
    )
    writer(
        "action_result",
        {
            "action_index": 1,
            "request_id": "action-1",
            "idempotent_replay": False,
            "receipt": {"ok": True},
            "screenshot": screenshot,
            "action": {"action": "done", "intent": "Finish."},
            "tool_use_id": "call-1",
            "batch_index": 1,
            "batch_size": 1,
            "batch_start_action_index": 1,
        },
    )
    writer.finish(
        {
            "completed": True,
            "stop_reason": "done",
            "failure": None,
            "accepted_actions": 1,
            "invalid_actions": 0,
            "screenshots": 1,
        }
    )

    assert (output_dir / "reference.png").read_bytes() == _PNG
    assert (output_dir / "frames" / "frame_0000.png").read_bytes() == _PNG
    turn = json.loads((output_dir / "model_turns.jsonl").read_text())
    assert turn["model_input"]["model"] == "gpt-5.6-sol"
    compact_turn = json.loads((output_dir / "model_turn_summaries.jsonl").read_text())
    assert compact_turn["model_input"] == {
        "model": "gpt-5.6-sol",
        "artifacts": [
            {
                "type": "image",
                "artifact_path": "frames/frame_0000.png",
                "sha256": "abc",
                "bytes": 123,
            }
        ],
    }
    assert "very large repeated history" not in json.dumps(compact_turn)
    receipt = json.loads((output_dir / "input_receipts.jsonl").read_text())
    assert receipt["receipt"] == {"ok": True}
    assert receipt["request_id"] == "action-1"
    assert receipt["idempotent_replay"] is False
    assert receipt["tool_use_id"] == "call-1"
    assert receipt["batch_index"] == 1
    assert receipt["batch_size"] == 1
    assert receipt["action"]["intent"] == "Finish."
    state = json.loads((output_dir / "run-state.json").read_text())
    assert state["state"] == "completed"
    assert state["accepted_actions"] == 1
    assert state["usage"] == {
        "input_tokens": 10,
        "output_tokens": 2,
        "total_tokens": 12,
    }
    assert not (output_dir / "run-state.json.tmp").exists()