tmux-shot/tests/test_agent.py

from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any

from tmux_shot import TmuxApp
from tmux_shot.agent import AgentRunner, _split_run_on_chunk


class FakeFunction:
    def __init__(self, name: str, arguments: str) -> None:
        self.name = name
        self.arguments = arguments


class FakeToolCall:
    def __init__(self, call_id: str, name: str, arguments: str) -> None:
        self.id = call_id
        self.function = FakeFunction(name, arguments)


class FakeMessage:
    def __init__(self, content: str | None = None, tool_calls: list[FakeToolCall] | None = None) -> None:
        self.content = content
        self.tool_calls = tool_calls or []

    def model_dump(self, exclude_none: bool = True) -> dict[str, Any]:
        d: dict[str, Any] = {"role": "assistant", "content": self.content}
        if self.tool_calls:
            d["tool_calls"] = [
                {
                    "id": tc.id,
                    "type": "function",
                    "function": {"name": tc.function.name, "arguments": tc.function.arguments},
                }
                for tc in self.tool_calls
            ]
        if exclude_none:
            d = {k: v for k, v in d.items() if v is not None}
        return d


class FakeCompletions:
    def __init__(self, script: list[FakeMessage]) -> None:
        self._script = list(script)
        self.requests: list[dict[str, Any]] = []

    def create(self, **kwargs: Any):
        # snapshot messages: it's the same mutable list AgentRunner keeps appending
        # to for the rest of the run, so store a copy or every recorded request
        # would end up reflecting the final state instead of the state at call time.
        self.requests.append({**kwargs, "messages": list(kwargs["messages"])})
        item = self._script.pop(0)
        if isinstance(item, Exception):
            raise item
        return type("R", (), {"choices": [type("C", (), {"message": item})()]})()


class FakeChat:
    def __init__(self, script: list[FakeMessage]) -> None:
        self.completions = FakeCompletions(script)


class FakeClient:
    def __init__(self, script: list[FakeMessage]) -> None:
        self.chat = FakeChat(script)


def test_agent_runner_drives_app_via_scripted_tool_calls(demo_app: TmuxApp, out_dir: Path) -> None:
    time.sleep(0.3)

    script = [
        FakeMessage(
            content="Moving down.",
            tool_calls=[FakeToolCall("call_1", "send_keys", '{"keys": "Down Down", "enter": false}')],
        ),
        FakeMessage(
            tool_calls=[FakeToolCall("call_2", "screenshot", '{"label": "moved"}')],
        ),
        FakeMessage(
            tool_calls=[FakeToolCall("call_3", "finish", '{"summary": "moved and captured"}')],
        ),
    ]
    client = FakeClient(script)

    runner = AgentRunner(demo_app, client, model="fake-model", out_dir=out_dir, delay=0.05, max_steps=5)
    runner.run("move down twice and screenshot it")

    assert len(runner.shots) == 1
    assert Path(runner.shots[0]).exists()

    # 3 request round-trips: send_keys turn, screenshot turn, finish turn
    assert len(client.chat.completions.requests) == 3

    # the request after the screenshot tool call must include the image as a follow-up
    # user message, immediately after that turn's tool-response message
    third_request_messages = client.chat.completions.requests[2]["messages"]
    roles = [m["role"] for m in third_request_messages]
    assert roles.count("tool") == 2  # one per completed tool call so far
    image_messages = [
        m for m in third_request_messages if m["role"] == "user" and isinstance(m.get("content"), list)
    ]
    assert len(image_messages) == 1
    content_types = [part["type"] for part in image_messages[0]["content"]]
    assert "image_url" in content_types


def test_agent_runner_stops_when_model_returns_no_tool_calls(demo_app: TmuxApp, out_dir: Path) -> None:
    script = [FakeMessage(content="Nothing to do here.")]
    client = FakeClient(script)

    runner = AgentRunner(demo_app, client, model="fake-model", out_dir=out_dir, delay=0.05, max_steps=5)
    runner.run("do nothing")

    assert runner.shots == []
    assert len(client.chat.completions.requests) == 1


def test_agent_runner_keeps_capturing_via_ticker_even_if_model_refuses_outright(
    demo_app: TmuxApp, out_dir: Path
) -> None:
    """Reproduces a real failure: the model flatly refuses a 'gif of N minutes of X' request
    ("I'm sorry, but I can't help with that.") instead of calling any tool. With --duration
    set, the background screenshot ticker must still cover the full requested wall-clock time
    regardless of the model's cooperation, so a timed GIF/video request is a mechanical
    guarantee, not something that depends on the model understanding the ask."""
    refusal = FakeMessage(content="I'm sorry, but I can't help with that.")
    client = FakeClient([refusal] * 10)  # more than enough for every nudge attempt

    runner = AgentRunner(
        demo_app,
        client,
        model="fake-model",
        out_dir=out_dir,
        delay=0.0,
        max_steps=50,
        duration=1.2,
        screenshot_interval=0.2,
    )

    start = time.monotonic()
    runner.run("I want a gif of you actively using this app")
    elapsed = time.monotonic() - start

    assert elapsed >= 1.0  # actually waited out roughly the full duration, not just bailed
    assert len(runner.shots) >= 3  # the ticker produced real frames despite total refusal
    for shot in runner.shots:
        assert Path(shot).exists()


def test_agent_runner_survives_malformed_tool_call_instead_of_crashing(
    demo_app: TmuxApp, out_dir: Path
) -> None:
    """Reproduces a real crash: the model called send_keys({}) -- no 'keys' argument at all --
    which used to raise an unhandled KeyError and kill the whole process. It must instead come
    back as a tool-result error the model can read and recover from."""
    script = [
        FakeMessage(tool_calls=[FakeToolCall("call_1", "send_keys", "{}")]),  # missing "keys"
        FakeMessage(tool_calls=[FakeToolCall("call_2", "send_keys", "not json")]),  # malformed JSON
        FakeMessage(tool_calls=[FakeToolCall("call_3", "finish", '{"summary": "recovered"}')]),
    ]
    client = FakeClient(script)

    runner = AgentRunner(demo_app, client, model="fake-model", out_dir=out_dir, delay=0.0, max_steps=5)
    runner.run("do something")  # must not raise

    third_request_messages = client.chat.completions.requests[2]["messages"]
    tool_messages = [m for m in third_request_messages if m["role"] == "tool"]
    assert len(tool_messages) == 2
    assert "error" in tool_messages[0]["content"]
    assert "missing required argument" in tool_messages[0]["content"]
    assert "error" in tool_messages[1]["content"]


def test_split_run_on_chunk_splits_bundled_numbered_items() -> None:
    bundled = (
        "1. Open vim in a new terminal and create a file. "
        "2. Use fzf to search project files and jump to a definition. "
        "3. Duplicate a heading with a macro and save."
    )
    parts = _split_run_on_chunk(bundled)
    assert len(parts) == 3
    assert parts[0].startswith("Open vim")
    assert parts[1].startswith("Use fzf")
    assert parts[2].startswith("Duplicate a heading")


def test_split_run_on_chunk_leaves_ordinary_text_alone() -> None:
    assert _split_run_on_chunk("create a macro with qa and replay it with 5@a") == [
        "create a macro with qa and replay it with 5@a"
    ]


def test_agent_runner_normalizes_a_bundled_plan_into_separate_chunks(
    demo_app: TmuxApp, out_dir: Path, capsys
) -> None:
    """Reproduces a real failure: the model returned a plan with all 3 steps bundled into a
    single array element ('1. ... 2. ... 3. ...') instead of 3 separate elements as instructed.
    The runner must recover this into real separate chunks rather than executing one giant,
    unbounded mega-chunk."""
    bundled_plan_call = FakeToolCall(
        "plan_1",
        "propose_plan",
        json.dumps(
            {
                "chunks": [
                    "1. Open vim and create a file. 2. Record a macro and replay it 5 times. "
                    "3. Open a vertical split and copy a line across."
                ]
            }
        ),
    )
    script = [
        FakeMessage(tool_calls=[bundled_plan_call]),
        FakeMessage(tool_calls=[FakeToolCall("c1", "finish", '{"summary": "opened file"}')]),
        FakeMessage(tool_calls=[FakeToolCall("c2", "finish", '{"summary": "did macro"}')]),
        FakeMessage(tool_calls=[FakeToolCall("c3", "finish", '{"summary": "did split"}')]),
    ]
    client = FakeClient(script)

    runner = AgentRunner(
        demo_app,
        client,
        model="fake-model",
        out_dir=out_dir,
        delay=0.0,
        max_steps=12,
        duration=6.0,
        screenshot_interval=100.0,
        chunk_seconds=2.0,
    )
    runner.run("give a demo of advanced vim features")

    out = capsys.readouterr().out
    assert "split into 3 chunks" in out
    assert "chunk 3/3 done: did split" in out


def test_agent_runner_chunks_a_long_duration_session_with_status_updates(
    demo_app: TmuxApp, out_dir: Path, capsys
) -> None:
    """A long --duration task should be planned into chunks (like an outline before drafting
    chapters) and each chunk run with its own fresh context, printing a status line per chunk
    as it completes."""
    plan_call = FakeToolCall(
        "plan_1", "propose_plan", '{"chunks": ["show text objects", "show macros", "show marks"]}'
    )
    script = [
        FakeMessage(tool_calls=[plan_call]),
        FakeMessage(tool_calls=[FakeToolCall("c1", "finish", '{"summary": "did text objects"}')]),
        FakeMessage(tool_calls=[FakeToolCall("c2", "finish", '{"summary": "did macros"}')]),
        FakeMessage(tool_calls=[FakeToolCall("c3", "finish", '{"summary": "did marks"}')]),
    ]
    client = FakeClient(script)

    runner = AgentRunner(
        demo_app,
        client,
        model="fake-model",
        out_dir=out_dir,
        delay=0.0,
        max_steps=12,
        duration=6.0,
        screenshot_interval=100.0,  # long enough the ticker won't fire during this fast fake run
        chunk_seconds=2.0,  # 6s / 2s -> plans for ~3 chunks
    )
    runner.run("give a demo of advanced vim features")

    assert len(client.chat.completions.requests) == 4  # 1 planning call + 3 chunk turns

    out = capsys.readouterr().out
    assert "plan (3 chunks)" in out
    assert "chunk 1/3 done: did text objects" in out
    assert "chunk 2/3 done: did macros" in out
    assert "chunk 3/3 done: did marks" in out


def test_agent_runner_retries_transient_api_errors_instead_of_crashing(
    demo_app: TmuxApp, out_dir: Path
) -> None:
    """Reproduces a real crash: the upstream gateway returned a transient 502
    (openai.InternalServerError) mid-session, which used to propagate straight up and kill the
    whole run. It must be retried with backoff instead."""
    import httpx2
    from openai import InternalServerError

    request = httpx2.Request("POST", "https://example.invalid/v1/chat/completions")
    response = httpx2.Response(502, request=request, json={"detail": "Upstream connection failed"})
    transient_error = InternalServerError("Bad Gateway", response=response, body=None)

    script = [
        transient_error,
        transient_error,
        FakeMessage(tool_calls=[FakeToolCall("call_1", "finish", '{"summary": "done after retry"}')]),
    ]
    client = FakeClient(script)

    runner = AgentRunner(demo_app, client, model="fake-model", out_dir=out_dir, delay=0.0, max_steps=5)
    runner.run("do something")  # must not raise

    assert len(client.chat.completions.requests) == 3


def test_agent_runner_degrades_gracefully_on_persistent_backend_outage(
    demo_app: TmuxApp, out_dir: Path, monkeypatch
) -> None:
    """Reproduces a real outage: the backend stayed down for longer than the retry budget
    (every call 502s). run() must not crash with a raw traceback -- it should stop cleanly and
    keep reporting whatever screenshots were already captured."""
    import httpx2
    from openai import InternalServerError

    import tmux_shot.agent as agent_module

    monkeypatch.setattr(agent_module.time, "sleep", lambda _seconds: None)  # skip real backoff waits

    request = httpx2.Request("POST", "https://example.invalid/v1/chat/completions")
    response = httpx2.Response(502, request=request, json={"detail": "Upstream connection failed"})
    always_down = InternalServerError("Bad Gateway", response=response, body=None)

    client = FakeClient(
        [
            FakeMessage(tool_calls=[FakeToolCall("call_1", "screenshot", '{"label": "before-outage"}')]),
            always_down,
            always_down,
            always_down,
            always_down,
        ]
    )

    runner = AgentRunner(demo_app, client, model="fake-model", out_dir=out_dir, delay=0.0, max_steps=5)
    runner.run("do something")  # must not raise

    assert len(runner.shots) == 1  # the screenshot taken before the outage is not lost
    assert Path(runner.shots[0]).exists()


def test_agent_step_strips_leaked_chat_template_tokens_from_tool_name(
    demo_app: TmuxApp, out_dir: Path
) -> None:
    """Reproduces a real artifact: a backend leaked raw harmony-style channel tokens into the
    function name itself ("screenshot<|channel|>commentary" instead of "screenshot"). It must
    still dispatch as the real tool on the first try, not fall through to 'unknown tool'."""
    script = [
        FakeMessage(tool_calls=[FakeToolCall("call_1", "read_text<|channel|>commentary", "{}")]),
    ]
    client = FakeClient(script)

    runner = AgentRunner(demo_app, client, model="fake-model", out_dir=out_dir, delay=0.0, max_steps=5)
    messages: list[dict] = []
    had_tools, finished, _ = runner._step(messages, "step 1")

    tool_message = next(m for m in messages if m["role"] == "tool")
    assert "unknown tool" not in tool_message["content"]
    assert "error" not in tool_message["content"]


def test_agent_runner_respects_max_steps_safety_cap(demo_app: TmuxApp, out_dir: Path) -> None:
    # every turn asks to wait, never calling finish -> must stop at max_steps, not hang forever
    script = [
        FakeMessage(tool_calls=[FakeToolCall(f"call_{i}", "wait", '{"seconds": 0.01}')]) for i in range(10)
    ]
    client = FakeClient(script)

    runner = AgentRunner(demo_app, client, model="fake-model", out_dir=out_dir, delay=0.0, max_steps=3)
    runner.run("wait forever")

    assert len(client.chat.completions.requests) == 3