371 lines
14 KiB
Python
371 lines
14 KiB
Python
|
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
|