loop.py

#!/usr/bin/env python3
"""The outer loop, kept deliberately dumb.

    while not converged:
        run the agent          (any CLI that reads a task on stdin and edits files)
        render + diff + blame  (evaluate.py)
        write feedback for the next run

No planner, no orchestration graph, no conversation memory. All state lives in
files inside the run directory, which makes every iteration inspectable and the
whole thing resumable and agent-agnostic. The agent supplies the intelligence;
the evaluator supplies the honesty.

    python loop.py examples/pricing-card/target.png
    python loop.py shot.png --agent "cursor-agent -p" --iters 10
"""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path

from PIL import Image

import evaluate

ROOT = Path(__file__).resolve().parent
DEFAULT_AGENT = "claude -p --dangerously-skip-permissions"
ARTIFACTS = ("feedback.md", "render.png", "diff.png", "side_by_side.png")


def build_prompt(threshold: float) -> str:
    return (
        (ROOT / "PROMPT.md").read_text()
        .replace("__THRESHOLD_PCT__", f"{threshold:.1%}")
        .replace("__PYTHON__", sys.executable)
        .replace("__EVALUATE__", str(ROOT / "evaluate.py"))
    )


def setup_run(target: Path, run_dir: Path | None, threshold: float, agent: str) -> Path:
    if run_dir is None:
        name = target.parent.name if target.stem == "target" else target.stem
        run_dir = ROOT / "runs" / f"{name}-{datetime.now():%m%d-%H%M%S}"
    run_dir.mkdir(parents=True, exist_ok=True)

    img = Image.open(target)
    if img.width > 1600:
        print(f"note: target is {img.width}px wide — retina screenshots are 2x; "
              "a 1x capture converges faster and at true sizes")
    img.convert("RGBA").save(run_dir / "target.png")

    (run_dir / "PROMPT.md").write_text(build_prompt(threshold))
    (run_dir / "harness.json").write_text(json.dumps({
        "target": str(target), "size": list(img.size),
        "threshold": threshold, "agent": agent,
        "created": datetime.now().isoformat(timespec="seconds"),
    }, indent=2))
    return run_dir


def run_agent(cmd: str, run_dir: Path, prompt: str, timeout: int) -> int | None:
    """One agent turn: task on stdin, cwd = run dir, output streamed through."""
    try:
        return subprocess.run(cmd, shell=True, cwd=run_dir, input=prompt,
                              text=True, timeout=timeout).returncode
    except subprocess.TimeoutExpired:
        print(f"  agent hit the {timeout}s timeout — evaluating whatever it left behind")
        return None


def archive(run_dir: Path, iteration: int) -> None:
    dest = run_dir / "feedback" / f"iter-{iteration:02d}"
    dest.mkdir(parents=True, exist_ok=True)
    for name in ARTIFACTS:
        src = run_dir / "feedback" / name
        if src.exists():
            shutil.copy(src, dest / name)


def main() -> None:
    ap = argparse.ArgumentParser(
        description="Loop a coding agent against a UI screenshot until the render matches.")
    ap.add_argument("target", help="screenshot to replicate (png/jpg, ideally 1x scale)")
    ap.add_argument("--agent", default=DEFAULT_AGENT,
                    help="shell command that reads a task on stdin and edits files in cwd "
                         f"(default: {DEFAULT_AGENT!r})")
    ap.add_argument("--iters", type=int, default=8, help="max iterations (default 8)")
    ap.add_argument("--threshold", type=float, default=evaluate.DEFAULT_CONVERGENCE,
                    help="pixel-match fraction that counts as done (default 0.995)")
    ap.add_argument("--run-dir", type=Path, default=None,
                    help="workspace to use; pass an existing one to resume it")
    ap.add_argument("--agent-timeout", type=int, default=1800,
                    help="seconds per agent turn (default 1800)")
    args = ap.parse_args()

    target = Path(args.target)
    if not target.exists():
        sys.exit(f"target not found: {target}")

    run_dir = setup_run(target, args.run_dir, args.threshold, args.agent)
    prompt = (run_dir / "PROMPT.md").read_text()
    history_path = run_dir / "feedback" / "history.jsonl"
    history = ([json.loads(l) for l in history_path.read_text().splitlines()]
               if history_path.exists() else [])
    start = len(history) + 1

    size = Image.open(run_dir / "target.png").size
    print(f"target    {target}  ({size[0]}x{size[1]})")
    print(f"run dir   {run_dir}")
    print(f"agent     {args.agent}")
    print(f"done at   {args.threshold:.1%} pixel match, max {args.iters} iterations")

    best = max((h["score"] for h in history), default=0.0)
    try:
        for i in range(start, args.iters + 1):
            print(f"\n--- iteration {i}/{args.iters} " + "-" * 40)
            t0 = time.time()
            rc = run_agent(args.agent, run_dir, prompt, args.agent_timeout)
            if rc:
                print(f"  agent exited with code {rc} — evaluating anyway")

            result = evaluate.evaluate(run_dir, threshold=args.threshold, iteration=i)
            if result is None:
                print("  no index.html produced yet")
                continue

            result["agent_seconds"] = round(time.time() - t0 - result["eval_seconds"], 1)
            history_path.parent.mkdir(exist_ok=True)
            with history_path.open("a") as f:
                f.write(json.dumps(result) + "\n")
            archive(run_dir, i)

            trend = "  (new best)" if result["score"] > best else ""
            best = max(best, result["score"])
            print(f"  match {result['score']:.2%}  ({result['mismatched']:,} px off, "
                  f"agent {result['agent_seconds']:.0f}s, eval {result['eval_seconds']}s){trend}")

            if result["converged"]:
                print(f"\nconverged: {result['score']:.2%} >= {args.threshold:.2%} "
                      f"after {i} iteration(s)")
                break
        else:
            print(f"\nstopped after {args.iters} iterations — best {best:.2%}")
    except KeyboardInterrupt:
        print("\ninterrupted — resume with:")
        print(f"  python loop.py {target} --run-dir {run_dir}")

    print(f"\nresult   {run_dir / 'index.html'}")
    print(f"compare  {run_dir / 'feedback' / 'side_by_side.png'}")


if __name__ == "__main__":
    main()